Support body rewrite

server/proxy.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"bytes"
 	"io"
 	"local/rproxy3/config"
 	"local/rproxy3/storage/packable"
@@ -23,7 +24,7 @@ type rewrite struct {
 }
 
 func (s *Server) Proxy(w http.ResponseWriter, r *http.Request) {
-	newURL, err := s.lookup(r.Host)
+	newURL, err := s.lookup(mapKey(r.Host))
 	var transport http.RoundTripper
 	transport = &redirPurge{
 		proxyHost:     r.Host,
@@ -31,7 +32,7 @@ func (s *Server) Proxy(w http.ResponseWriter, r *http.Request) {
 		baseTransport: http.DefaultTransport,
 	}
 	transport = &rewrite{
-		rewrites:      config.GetRewrites(),
+		rewrites:      config.GetRewrites(mapKey(r.Host)),
 		baseTransport: transport,
 	}
 	if err != nil {
@@ -46,13 +47,17 @@ func (s *Server) Proxy(w http.ResponseWriter, r *http.Request) {
 }
 
 func (s *Server) lookup(host string) (*url.URL, error) {
-	host = strings.Split(host, ".")[0]
-	host = strings.Split(host, ":")[0]
 	v := packable.NewURL()
 	err := s.db.Get(nsRouting, host, v)
 	return v.URL(), err
 }
 
+func mapKey(host string) string {
+	host = strings.Split(host, ".")[0]
+	host = strings.Split(host, ":")[0]
+	return host
+}
+
 func (rp *redirPurge) RoundTrip(r *http.Request) (*http.Response, error) {
 	resp, err := rp.baseTransport.RoundTrip(r)
 	if err != nil {
@@ -69,13 +74,22 @@ func (rw *rewrite) RoundTrip(r *http.Request) (*http.Response, error) {
 	if err != nil {
 		return resp, err
 	}
+	if len(rw.rewrites) == 0 {
+		return resp, err
+	}
+	resp.Header.Del("Content-Length")
 	pr, pw := io.Pipe()
 	body := resp.Body
 	resp.Body = pr
 	go func() {
 		buff := make([]byte, 1024)
-		for n, err := body.Read(buff); err == nil || n > 0; n, err = body.Read(buff) {
+		n, err := body.Read(buff)
+		for err == nil || n > 0 {
 			chunk := buff[:n]
+			for k, v := range rw.rewrites {
+				chunk = bytes.Replace(chunk, []byte(k), []byte(v), -1)
+			}
+			n = len(chunk)
 			m := 0
 			for m < n {
 				l, err := pw.Write(chunk[m:])
@@ -85,6 +99,7 @@ func (rw *rewrite) RoundTrip(r *http.Request) (*http.Response, error) {
 				}
 				m += l
 			}
+			n, err = body.Read(buff)
 		}
 		pw.CloseWithError(err)
 	}()