add response filter

Author: phuslu

direct.go

@@ -1,16 +1,21 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
 )
 
-type DirectPlugin struct {
-	Plugin
+type DirectRequestPlugin struct {
+	RequestPlugin
 }
 
-func (p DirectPlugin) Handle(c *PluginContext, rw http.ResponseWriter, req *http.Request) {
+type DirectResponsePlugin struct {
+	ResponsePlugin
+}
+
+func (p DirectRequestPlugin) HandleRequest(c *PluginContext, rw http.ResponseWriter, req *http.Request) (*http.Response, error) {
 	if req.Method != "CONNECT" {
 		if !req.URL.IsAbs() {
 			if req.TLS != nil {
@@ -29,14 +34,30 @@ func (p DirectPlugin) Handle(c *PluginContext, rw http.ResponseWriter, req *http
 		if err != nil {
 			rw.WriteHeader(502)
 			fmt.Fprintf(rw, "Error: %s\n", err)
-			return
+			return nil, err
 		}
 		newReq.Header = req.Header
 		res, err := c.H.Net.HttpClientDo(newReq)
-		if err != nil {
+		return res, err
+	} else {
+		c.H.Log.Printf("%s \"DIRECT %s %s %s\" - -", req.RemoteAddr, req.Method, req.Host, req.Proto)
+		response := &http.Response{
+			StatusCode:    200,
+			ProtoMajor:    1,
+			ProtoMinor:    1,
+			Header:        http.Header{},
+			ContentLength: -1,
+		}
+		return response, nil
+	}
+}
+
+func (p DirectResponsePlugin) HandleResponse(c *PluginContext, rw http.ResponseWriter, req *http.Request, res *http.Response, resError error) error {
+	if req.Method != "CONNECT" {
+		if resError != nil {
 			rw.WriteHeader(502)
-			fmt.Fprintf(rw, "Error: %s\n", err)
-			return
+			fmt.Fprintf(rw, "Error: %s\n", resError)
+			return resError
 		}
 		c.H.Log.Printf("%s \"DIRECT %s %s %s\" %d %s", req.RemoteAddr, req.Method, req.URL.String(), req.Proto, res.StatusCode, res.Header.Get("Content-Length"))
 		rw.WriteHeader(res.StatusCode)
@@ -47,20 +68,27 @@ func (p DirectPlugin) Handle(c *PluginContext, rw http.ResponseWriter, req *http
 		}
 		io.Copy(rw, res.Body)
 	} else {
-		c.H.Log.Printf("%s \"DIRECT %s %s %s\" - -", req.RemoteAddr, req.Method, req.Host, req.Proto)
+		if resError != nil {
+			rw.WriteHeader(502)
+			fmt.Fprintf(rw, "Error: %s\n", resError)
+			c.H.Log.Printf("NetDialTimeout %s failed %s", req.Host, resError)
+			return resError
+		}
 		remoteConn, err := c.H.Net.NetDialTimeout("tcp", req.Host, c.H.Net.GetTimeout())
 		if err != nil {
-			c.H.Log.Printf("NetDialTimeout %s failed %s", req.Host, err)
-			return
+			return err
 		}
 		hijacker, ok := rw.(http.Hijacker)
 		if !ok {
-			c.H.Log.Printf("http.ResponseWriter does not implments Hijacker")
-			return
+			resError = errors.New("http.ResponseWriter does not implments Hijacker")
+			rw.WriteHeader(502)
+			fmt.Fprintf(rw, "Error: %s\n", resError)
+			return resError
 		}
 		localConn, _, err := hijacker.Hijack()
 		localConn.Write([]byte("HTTP/1.1 200 OK\r\n\r\n"))
 		go io.Copy(remoteConn, localConn)
 		io.Copy(localConn, remoteConn)
 	}
+	return nil
 }

filter.go

@@ -4,19 +4,27 @@ import (
 	"net/http"
 )
 
-type DirectFilter struct {
+type DirectRequestFilter struct {
 	RequestFilter
 }
 
-func (d *DirectFilter) Filter(req *http.Request) (pluginName string, pluginArgs *http.Header, err error) {
+func (d *DirectRequestFilter) Filter(req *http.Request) (pluginName string, pluginArgs *http.Header, err error) {
 	return "direct", nil, nil
 }
 
-type StripFilter struct {
+type DirectResponseFilter struct {
 	RequestFilter
 }
 
-func (d *StripFilter) Filter(req *http.Request) (pluginName string, pluginArgs *http.Header, err error) {
+func (d *DirectResponseFilter) Filter(req *http.Request, res *http.Response) (pluginName string, pluginArgs *http.Header, err error) {
+	return "direct", nil, nil
+}
+
+type StripRequestFilter struct {
+	RequestFilter
+}
+
+func (d *StripRequestFilter) Filter(req *http.Request) (pluginName string, pluginArgs *http.Header, err error) {
 	if req.Method == "CONNECT" {
 		args := http.Header{
 			"Foo": []string{"bar"},

goagent.go

@@ -20,13 +20,19 @@ func main() {
 		Listener: ln,
 		Log:      log.New(os.Stderr, "INFO - ", 3),
 		Net:      &SimpleNetwork{},
-		Plugins: map[string]Plugin{
-			"direct": &DirectPlugin{},
-			"strip":  &StripPlugin{},
+		RequestPlugins: map[string]RequestPlugin{
+			"direct": &DirectRequestPlugin{},
+			"strip":  &StripRequestPlugin{},
+		},
+		ResponsePlugins: map[string]ResponsePlugin{
+			"direct": &DirectResponsePlugin{},
 		},
 		RequestFilters: []RequestFilter{
-			&StripFilter{},
-			&DirectFilter{},
+			&StripRequestFilter{},
+			&DirectRequestFilter{},
+		},
+		ResponseFilters: []ResponseFilter{
+			&DirectResponseFilter{},
 		},
 	}
 	s := &http.Server{

handler.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"crypto/tls"
+	"io"
 	"log"
 	"net"
 	"net/http"
@@ -13,6 +14,7 @@ type Net2 interface {
 	NetDialTimeout(network, address string, timeout time.Duration) (net.Conn, error)
 	TlsDialTimeout(network, address string, config *tls.Config, timeout time.Duration) (*tls.Conn, error)
 	HttpClientDo(req *http.Request) (*http.Response, error)
+	CopyResponseBody(w io.Writer, res *http.Response) (int64, error)
 	GetTimeout() time.Duration
 	SetTimeout()
 	GetAddressAlias(addr string) (alias string)
@@ -23,7 +25,7 @@ type RequestFilter interface {
 }
 
 type ResponseFilter interface {
-	Filter(req *http.Response) (newReq *http.Response, err error)
+	Filter(req *http.Request, res *http.Response) (pluginName string, pluginArgs *http.Header, err error)
 }
 
 type PushListener interface {
@@ -100,7 +102,8 @@ type Handler struct {
 	Listener        net.Listener
 	Log             *log.Logger
 	Net             Net2
-	Plugins         map[string]Plugin
+	RequestPlugins  map[string]RequestPlugin
+	ResponsePlugins map[string]ResponsePlugin
 	RequestFilters  []RequestFilter
 	ResponseFilters []ResponseFilter
 }
@@ -110,22 +113,53 @@ type PluginContext struct {
 	Args *http.Header
 }
 
+type RequestPlugin interface {
+	HandleRequest(*PluginContext, http.ResponseWriter, *http.Request) (*http.Response, error)
+}
+
+type ResponsePlugin interface {
+	HandleResponse(*PluginContext, http.ResponseWriter, *http.Request, *http.Response, error) error
+}
+
 type Plugin interface {
-	Handle(*PluginContext, http.ResponseWriter, *http.Request)
+	RequestPlugin
+	ResponsePlugin
 }
 
 func (h Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
-	for _, f := range h.RequestFilters {
-		name, args, err := f.Filter(req)
+	for _, reqfilter := range h.RequestFilters {
+		name, args, err := reqfilter.Filter(req)
 		if err != nil {
-			h.Log.Fatalf("ServeHTTP error: %v", err)
+			h.Log.Printf("RequestFilter error: %v", err)
 		}
 		if name == "" {
 			continue
 		}
-		if plugin, ok := h.Plugins[name]; ok {
-			context := &PluginContext{&h, args}
-			plugin.Handle(context, rw, req)
+		if reqplugin, ok := h.RequestPlugins[name]; ok {
+			reqctx := &PluginContext{&h, args}
+			res, err := reqplugin.HandleRequest(reqctx, rw, req)
+			if err != nil {
+				h.Log.Printf("Plugin %s HandleResponse error: %v", name, err)
+			}
+			if res != nil {
+				for _, resfilter := range h.ResponseFilters {
+					name, args, err := resfilter.Filter(req, res)
+					if err != nil {
+						h.Log.Printf("ServeHTTP RequestFilter error: %v", err)
+					}
+					if name == "" {
+						continue
+					}
+					if resplugin, ok := h.ResponsePlugins[name]; ok {
+						resctx := &PluginContext{&h, args}
+						err := resplugin.HandleResponse(resctx, rw, req, res, err)
+						if err != nil {
+							h.Log.Printf("Plugin %s HandleResponse error: %v", name, err)
+						}
+					}
+					break
+				}
+			}
 			break
 		} else {
 			h.Log.Fatalf("plugin \"%s\" not registered", name)

network.go

@@ -2,39 +2,86 @@ package main
 
 import (
 	"crypto/tls"
+	// "fmt"
+	"io"
 	"net"
 	"net/http"
+	"sync"
 	"time"
 )
 
 type SimpleNetwork struct {
 	Net2
 }
 
-func (s *SimpleNetwork) NetResolveIPAddr(network, addr string) (*net.IPAddr, error) {
+func (sn *SimpleNetwork) NetResolveIPAddr(network, addr string) (*net.IPAddr, error) {
 	return net.ResolveIPAddr(network, addr)
 }
 
-func (s *SimpleNetwork) NetDialTimeout(network, address string, timeout time.Duration) (net.Conn, error) {
+func (sn *SimpleNetwork) NetDialTimeout(network, address string, timeout time.Duration) (net.Conn, error) {
 	return net.DialTimeout(network, address, timeout)
 }
 
-func (s *SimpleNetwork) TlsDialTimeout(network string, addr string, config *tls.Config, timeout time.Duration) (*tls.Conn, error) {
+func (sn *SimpleNetwork) TlsDialTimeout(network string, addr string, config *tls.Config, timeout time.Duration) (*tls.Conn, error) {
 	return tls.Dial(network, addr, config)
 }
 
-func (s *SimpleNetwork) HttpClientDo(req *http.Request) (*http.Response, error) {
+func (sn *SimpleNetwork) HttpClientDo(req *http.Request) (*http.Response, error) {
 	client := &http.Client{}
 	return client.Do(req)
 }
 
-func (s *SimpleNetwork) GetTimeout() time.Duration {
+func (sn *SimpleNetwork) CopyResponseBody(w io.Writer, res *http.Response) (int64, error) {
+	return io.Copy(w, res.Body)
+}
+
+func (sn *SimpleNetwork) GetTimeout() time.Duration {
+	return 8 * time.Second
+}
+
+func (sn *SimpleNetwork) SetTimeout() {
+}
+
+func (sn *SimpleNetwork) GetAddressAlias(addr string) (alias string) {
+	return ""
+}
+
+type AdvancedNetwork struct {
+	Net2
+	dnsCache   map[string]*net.IPAddr
+	dnsCacheMu sync.Mutex
+}
+
+func NewAdvancedNetwork() *AdvancedNetwork {
+	return &AdvancedNetwork{
+		dnsCache: map[string]*net.IPAddr{},
+	}
+}
+
+func (an *AdvancedNetwork) NetResolveIPAddr(network, addr string) (*net.IPAddr, error) {
+	return net.ResolveIPAddr(network, addr)
+}
+
+func (an *AdvancedNetwork) NetDialTimeout(network, address string, timeout time.Duration) (net.Conn, error) {
+	return net.DialTimeout(network, address, timeout)
+}
+
+func (an *AdvancedNetwork) TlsDialTimeout(network string, addr string, config *tls.Config, timeout time.Duration) (*tls.Conn, error) {
+	return tls.Dial(network, addr, config)
+}
+
+func (an *AdvancedNetwork) HttpClientDo(req *http.Request) (*http.Response, error) {
+	client := &http.Client{}
+	return client.Do(req)
+}
+
+func (an *AdvancedNetwork) GetTimeout() time.Duration {
 	return 8 * time.Second
 }
 
-func (s *SimpleNetwork) SetTimeout() {
+func (s *AdvancedNetwork) SetTimeout() {
 }
 
-func (s *SimpleNetwork) GetAddressAlias(addr string) (alias string) {
+func (s *AdvancedNetwork) GetAddressAlias(addr string) (alias string) {
 	return ""
 }