(Go) Modify websocket body in reverse proxy - go

Question

Hello Everyone, I'm new to Go.
I'm creating a reverse proxy server using Go.
My Server has websocket. I finally get it connected.
Now I want to change websocket message body.
Sorry If my code is weird to you. Forgive me, I'm new to Go 😢
I'm wraping resp.body to NewReadWriteBody() in which contains wrapper for Read, Write and Closer. And I'm modifying message body inside it.
Here is how I'm doing with it:
package rever
// https://blog.joshsoftware.com/2021/05/25/simple-and-powerful-reverseproxy-in-go/
// https://github.com/golang/go/blob/master/src/net/http/httputil/reverseproxy.go
import (
"bytes"
"crypto/tls"
"fmt"
"io"
"io/ioutil"
"log"
"net/http"
"net/http/httputil"
"net/url"
"strconv"
)
type ReadWriteBody struct {
originBody io.ReadWriteCloser
}
func NewReadWriteBody(body io.ReadCloser) *ReadWriteBody {
b := &ReadWriteBody{}
rw, ok := body.(io.ReadWriteCloser)
if !ok {
log.Println("29: error while casting body to ReadWriteCloser")
}
b.originBody = rw
return b
}
func (b *ReadWriteBody) Read(p []byte) (n int, err error) {
buf := make([]byte, len(p))
n, err = b.originBody.Read(buf)
if err != nil {
log.Println("43: ", err.Error())
return n, err
}
buf = bytes.ReplaceAll(buf, []byte("mm.remote"), []byte("mm.local"))
copy(p[:], buf)
return len(p), nil
}
func (b *ReadWriteBody) Write(p []byte) (n int, err error) {
buf := make([]byte, len(p))
n, err = b.originBody.Write(buf)
if err != nil {
log.Println(err.Error())
return n, err
}
buf = bytes.ReplaceAll(buf, []byte("mm.local"), []byte("mm.remote"))
copy(p[:], buf)
return len(p), nil
}
func (b *ReadWriteBody) Close() error {
return b.originBody.Close()
}
type transport struct {
http.RoundTripper
}
func (t *transport) RoundTrip(req *http.Request) (resp *http.Response, err error) {
resp, err = t.RoundTripper.RoundTrip(req)
if err != nil {
log.Println("99: ", err.Error())
return nil, err
}
if resp.StatusCode == http.StatusSwitchingProtocols {
resp.Body = NewReadWriteBody(resp.Body)
return resp, nil
}
b, err := ioutil.ReadAll(resp.Body)
if err != nil {
log.Println("114: ", err.Error())
return nil, err
}
err = resp.Body.Close()
if err != nil {
log.Println("119", err.Error())
return nil, err
}
b = bytes.ReplaceAll(b, []byte("mm.remote"), []byte("mm.local"))
body := ioutil.NopCloser(bytes.NewReader(b))
resp.Body = body
resp.ContentLength = int64(len(b))
resp.Header.Set("Content-Length", strconv.Itoa(len(b)))
return resp, nil
}
var _ http.RoundTripper = &transport{}
// NewProxy takes target host and creates a reverse proxy
func NewProxy(targetHost string) (*httputil.ReverseProxy, error) {
url, err := url.Parse(targetHost)
if err != nil {
log.Println("141: ", err.Error())
return nil, err
}
proxy := httputil.NewSingleHostReverseProxy(url)
originalDirector := proxy.Director
proxy.Director = func(req *http.Request) {
originalDirector(req)
modifyRequest(req)
}
proxy.ErrorHandler = errorHandler()
dt := http.DefaultTransport.(*http.Transport).Clone()
dt.TLSClientConfig = &tls.Config{}
dt.ForceAttemptHTTP2 = false
proxy.Transport = &transport{dt}
return proxy, nil
}
func modifyRequest(req *http.Request) {
req.Host = "mm.remote"
req.Header.Set("Accept-Encoding", "identity")
}
func errorHandler() func(http.ResponseWriter, *http.Request, error) {
return func(w http.ResponseWriter, req *http.Request, err error) {
// fmt.Printf("Got error while modifying response: %v \n", err)
}
}
// ProxyRequestHandler handles the http request using proxy
func ProxyRequestHandler(proxy *httputil.ReverseProxy) func(http.ResponseWriter, *http.Request) {
return func(w http.ResponseWriter, r *http.Request) {
proxy.ServeHTTP(w, r)
}
}
func Main() {
// initialize a reverse proxy and pass the actual backend server url here
proxy, err := NewProxy("https://mm.remote")
if err != nil {
log.Println(err.Error())
panic(err)
}
// handle all requests to your server using the proxy
http.HandleFunc("/", ProxyRequestHandler(proxy))
fmt.Println("Server started")
log.Fatal(http.ListenAndServe(":8008", nil))
}