diff --git a/go.mod b/go.mod index f8ef008..496dca2 100644 --- a/go.mod +++ b/go.mod @@ -5,4 +5,7 @@ go 1.21.1 require ( github.com/andybalholm/brotli v1.0.6 github.com/gorilla/websocket v1.5.0 + github.com/sirupsen/logrus v1.9.3 ) + +require golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 // indirect diff --git a/go.sum b/go.sum index 93968f7..b1504da 100644 --- a/go.sum +++ b/go.sum @@ -1,4 +1,19 @@ github.com/andybalholm/brotli v1.0.6 h1:Yf9fFpf49Zrxb9NlQaluyE92/+X7UVHlhMNJN2sxfOI= github.com/andybalholm/brotli v1.0.6/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc= github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= +github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 h1:0A+M6Uqn+Eje4kHMK80dtF3JCXC4ykBgQG4Fe06QRhQ= +golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/main.go b/main.go index 1dca873..6626abd 100644 --- a/main.go +++ b/main.go @@ -5,25 +5,37 @@ import ( "compress/flate" "compress/gzip" "io" - "log" "net/http" "os" "strings" "github.com/andybalholm/brotli" "github.com/gorilla/websocket" -) - -var ( - debugLog *log.Logger - infoLog *log.Logger - errorLog *log.Logger + "github.com/sirupsen/logrus" ) func init() { - debugLog = log.New(os.Stdout, "DEBUG: ", log.Ldate|log.Ltime|log.Lshortfile) - infoLog = log.New(os.Stdout, "INFO: ", log.Ldate|log.Ltime|log.Lshortfile) - errorLog = log.New(os.Stderr, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile) + logLevel := os.Getenv("LOG_LEVEL") + if logLevel == "" { + logLevel = "info" + } + + switch logLevel { + case "debug": + logrus.SetLevel(logrus.DebugLevel) + case "info": + logrus.SetLevel(logrus.InfoLevel) + case "warning": + logrus.SetLevel(logrus.WarnLevel) + case "error": + logrus.SetLevel(logrus.ErrorLevel) + default: + logrus.SetLevel(logrus.InfoLevel) + } + + logrus.SetFormatter(&logrus.TextFormatter{ + FullTimestamp: true, + }) } func replaceDomainInResponse(originalSubdomain, replaceSubdomain, originalDomain string, buffer *bytes.Buffer) { @@ -40,14 +52,14 @@ func proxyRequest(fullSubdomain, path string, buffer *bytes.Buffer, r *http.Requ target := "https://" + fullSubdomain + ".lunaroasis.net" + path newReq, err := http.NewRequest(r.Method, target, r.Body) if err != nil { - errorLog.Printf("Failed to create request: %v", err) + logrus.Errorf("Failed to create request: %v", err) return 0, nil, err } newReq.Header = r.Header resp, err := client.Do(newReq) if err != nil { - errorLog.Printf("Failed to send request: %v", err) + logrus.Errorf("Failed to send request: %v", err) return 0, nil, err } defer resp.Body.Close() @@ -69,7 +81,7 @@ func proxyRequest(fullSubdomain, path string, buffer *bytes.Buffer, r *http.Requ // Decompress Gzip data reader, err = gzip.NewReader(resp.Body) if err != nil { - errorLog.Printf("Failed to create gzip reader: %v", err) + logrus.Errorf("Failed to create gzip reader: %v", err) return 0, nil, err } case "deflate": @@ -84,11 +96,11 @@ func proxyRequest(fullSubdomain, path string, buffer *bytes.Buffer, r *http.Requ } func handleHttpRequest(w http.ResponseWriter, r *http.Request) { - infoLog.Printf("Received request from %s", r.Host) + logrus.Infof("Received request from %s", r.Host) hostParts := strings.Split(r.Host, ".") if len(hostParts) < 3 { - errorLog.Printf("Invalid domain: %s", r.Host) + logrus.Errorf("Invalid domain: %s", r.Host) http.Error(w, "Invalid domain", http.StatusBadRequest) return } @@ -106,13 +118,13 @@ func handleHttpRequest(w http.ResponseWriter, r *http.Request) { buffer := new(bytes.Buffer) backupBuffer := new(bytes.Buffer) - debugLog.Printf("Proxying request to %s", subdomain+"-statescale") + logrus.Debugf("Proxying request to %s", subdomain+"-statescale") statusCode, headers, err := proxyRequest(subdomain+"-statescale", r.RequestURI, buffer, r) - debugLog.Printf("Received status code %d", statusCode) + logrus.Debugf("Received status code %d", statusCode) if err != nil || statusCode >= 400 { - debugLog.Printf("Proxying request to %s", subdomain+"-snapscale") + logrus.Debugf("Proxying request to %s", subdomain+"-snapscale") backupStatusCode, backupHeaders, _ := proxyRequest(subdomain+"-snapscale", r.RequestURI, backupBuffer, r) - debugLog.Printf("Received status code %d", backupStatusCode) + logrus.Debugf("Received status code %d", backupStatusCode) replaceDomainInResponse(subdomain, subdomain+"-snapscale", originalDomain, backupBuffer) @@ -152,12 +164,12 @@ func proxyWebSocketRequest(subdomain string, w http.ResponseWriter, r *http.Requ dialer := websocket.Dialer{} targetConn, resp, err := dialer.Dial(target, nil) if err != nil { - errorLog.Printf("Failed to connect to target: %v", err) + logrus.Errorf("Failed to connect to target: %v", err) if resp != nil { - errorLog.Printf("Handshake response status: %s", resp.Status) + logrus.Errorf("Handshake response status: %s", resp.Status) // Log all response headers for debugging for k, v := range resp.Header { - errorLog.Printf("%s: %s", k, v) + logrus.Errorf("%s: %s", k, v) } } http.Error(w, "Internal server error", http.StatusInternalServerError) @@ -168,7 +180,7 @@ func proxyWebSocketRequest(subdomain string, w http.ResponseWriter, r *http.Requ // Upgrade the client connection to a WebSocket connection clientConn, err := upgrader.Upgrade(w, r, nil) if err != nil { - errorLog.Printf("Failed to upgrade client connection: %v", err) + logrus.Errorf("Failed to upgrade client connection: %v", err) return // No need to send an error response, Upgrade already did if there was an error } defer clientConn.Close() @@ -178,12 +190,12 @@ func proxyWebSocketRequest(subdomain string, w http.ResponseWriter, r *http.Requ for { messageType, message, err := targetConn.ReadMessage() if err != nil { - errorLog.Printf("Failed to read from target: %v", err) + logrus.Errorf("Failed to read from target: %v", err) return } err = clientConn.WriteMessage(messageType, message) if err != nil { - errorLog.Printf("Failed to write to client: %v", err) + logrus.Errorf("Failed to write to client: %v", err) return } } @@ -192,12 +204,12 @@ func proxyWebSocketRequest(subdomain string, w http.ResponseWriter, r *http.Requ for { messageType, message, err := clientConn.ReadMessage() if err != nil { - errorLog.Printf("Failed to read from client: %v", err) + logrus.Errorf("Failed to read from client: %v", err) return } err = targetConn.WriteMessage(messageType, message) if err != nil { - errorLog.Printf("Failed to write to target: %v", err) + logrus.Errorf("Failed to write to target: %v", err) return } } @@ -230,7 +242,7 @@ func handleRequest(w http.ResponseWriter, r *http.Request) { } func main() { - infoLog.Println("Starting server on :8080") + logrus.Info("Starting server on :8080") http.HandleFunc("/", handleRequest) http.ListenAndServe(":8080", nil) }