From 30339d1d26eeaf732cbbd379c79a2e5496296de5 Mon Sep 17 00:00:00 2001
From: Jakub <jakub@cuth.sk>
Date: Tue, 30 Jan 2024 17:49:46 +0100
Subject: [PATCH] feat(GODT-1602): add more logs to proxy requests.

---
 server/proxy.go | 43 ++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 38 insertions(+), 5 deletions(-)

diff --git a/server/proxy.go b/server/proxy.go
index d8f46c0..1b50b20 100644
--- a/server/proxy.go
+++ b/server/proxy.go
@@ -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)
 	}
@@ -24,14 +26,42 @@ 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 {
+		body, err := io.ReadAll(r.Body)
+		if err != nil {
+			body = []byte(fmt.Sprintf(`{"read_error": "%v"}`, err))
+		}
+
+		fmt.Printf("PROXY RESP %d CHANGING TO 429: %s %s\n=====BODY=====\n%s\n", r.StatusCode, r.Request.Method, r.Request.RequestURI, string(body))
+		r.StatusCode = http.StatusTooManyRequests
+		r.Body = io.NopCloser(bytes.NewReader(body))
+	}
+	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)
@@ -140,6 +170,8 @@ type proxyServer struct {
 	origin, base string
 
 	transport http.RoundTripper
+
+	debug bool
 }
 
 func newProxyServer(origin, base string, transport http.RoundTripper) *proxyServer {
@@ -148,6 +180,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") != "",
 	}
 }
 
@@ -161,7 +194,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") {