Skip to content

Commit

Permalink
feat(GODT-1602): add more logs to proxy requests.
Browse files Browse the repository at this point in the history
  • Loading branch information
cuthix committed Feb 2, 2024
1 parent e8aa47d commit 7033168
Showing 1 changed file with 32 additions and 5 deletions.
37 changes: 32 additions & 5 deletions server/proxy.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,18 +4,20 @@ import (
"bytes"
"compress/gzip"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httputil"
"net/url"
"os"
"strings"

"github.com/ProtonMail/go-proton-api"
"github.com/gin-gonic/gin"
)

func newProxy(proxyOrigin, base, path string, transport http.RoundTripper) http.HandlerFunc {
origin, err := url.Parse(proxyOrigin)
func (s *proxyServer) newProxy(path string) http.HandlerFunc {
origin, err := url.Parse(s.origin)
if err != nil {
panic(err)
}
Expand All @@ -24,14 +26,36 @@ func newProxy(proxyOrigin, base, path string, transport http.RoundTripper) http.
Director: func(req *http.Request) {
req.URL.Scheme = origin.Scheme
req.URL.Host = origin.Host
req.URL.Path = origin.Path + strings.TrimPrefix(path, base)
req.URL.Path = origin.Path + strings.TrimPrefix(path, s.base)
req.Host = origin.Host
},

Transport: transport,
ModifyResponse: s.targetResponseModifier,
ErrorHandler: s.targetErrorHandler,

Transport: s.transport,
}).ServeHTTP
}

func (s *proxyServer) targetResponseModifier(r *http.Response) error {
if s.debug {
fmt.Printf("PROXY RESP %d: %s %s\n", r.StatusCode, r.Request.Method, r.Request.RequestURI)
}
if r.StatusCode == http.StatusInternalServerError || r.StatusCode == http.StatusBadGateway {
fmt.Printf("PROXY RESP %d CHANGING TO 429: %s %s\n", r.StatusCode, r.Request.Method, r.Request.RequestURI)
r.StatusCode = http.StatusTooManyRequests
}
return nil
}

func (s *proxyServer) targetErrorHandler(rw http.ResponseWriter, r *http.Request, targetError error) {
if s.debug {
fmt.Printf("PROXY ERROR: %s %s: %v\n", r.Method, r.RequestURI, targetError)
}

rw.WriteHeader(http.StatusTooManyRequests)
}

func (s *Server) handleProxy(base string) gin.HandlerFunc {
return func(c *gin.Context) {
proxy := newProxyServer(s.proxyOrigin, base, s.proxyTransport)
Expand Down Expand Up @@ -140,6 +164,8 @@ type proxyServer struct {
origin, base string

transport http.RoundTripper

debug bool
}

func newProxyServer(origin, base string, transport http.RoundTripper) *proxyServer {
Expand All @@ -148,6 +174,7 @@ func newProxyServer(origin, base string, transport http.RoundTripper) *proxyServ
origin: origin,
base: base,
transport: transport,
debug: os.Getenv("GO_PROTON_API_SERVER_LOGGER_ENABLED") != "",
}
}

Expand All @@ -161,7 +188,7 @@ func (s *proxyServer) handle(path string, h func(func(string) HandlerFunc) http.
buf := new(bytes.Buffer)

// Call the proxy, capturing whatever data it writes.
newProxy(s.origin, s.base, path, s.transport)(&writerWrapper{w, buf}, r)
s.newProxy(path)(&writerWrapper{w, buf}, r)

// If there is a gzip header entry, decode it.
if strings.Contains(w.Header().Get("Content-Encoding"), "gzip") {
Expand Down

0 comments on commit 7033168

Please sign in to comment.