From 20a364170a4ac4983f3485c965f8dffc8392d2da Mon Sep 17 00:00:00 2001 From: fangdingjun Date: Tue, 10 Jul 2018 09:33:28 +0800 Subject: [PATCH] move responseWriter to handler.go --- handler.go | 104 +++++++++++++++++++++++++++++++++++++++++++++++++++++ routers.go | 103 ---------------------------------------------------- 2 files changed, 104 insertions(+), 103 deletions(-) diff --git a/handler.go b/handler.go index 15a6724..9aa7578 100644 --- a/handler.go +++ b/handler.go @@ -1,6 +1,8 @@ package main import ( + "bufio" + "bytes" "fmt" "io" "log" @@ -9,7 +11,9 @@ import ( "strings" "time" + "github.com/fangdingjun/gnutls" auth "github.com/fangdingjun/go-http-auth" + nghttp2 "github.com/fangdingjun/nghttp2-go" ) // handler process the proxy request first(if enabled) @@ -231,3 +235,103 @@ func pipeAndClose(r1, r2 io.ReadWriteCloser) { <-ch } + +func handleHTTPClient(c net.Conn, handler http.Handler) { + tlsconn := c.(*gnutls.Conn) + if err := tlsconn.Handshake(); err != nil { + log.Println(err) + return + } + state := tlsconn.ConnectionState() + if state.NegotiatedProtocol == "h2" { + h2conn, err := nghttp2.Server(tlsconn, handler) + if err != nil { + log.Println(err) + } + h2conn.Run() + return + } + + defer c.Close() + r := bufio.NewReader(tlsconn) + buf := new(bytes.Buffer) + for { + req, err := http.ReadRequest(r) + if err != nil { + return + } + addr := tlsconn.RemoteAddr().String() + req.RemoteAddr = addr + rh := &responseHandler{ + c: tlsconn, + header: http.Header{}, + buf: buf, + } + handler.ServeHTTP(rh, req) + rh.Write(nil) + rh.buf.WriteTo(rh.c) + } +} + +type responseHandler struct { + c net.Conn + statusCode int + header http.Header + responseSend bool + w io.Writer + buf *bytes.Buffer +} + +func (r *responseHandler) WriteHeader(statusCode int) { + if r.responseSend { + return + } + r.buf.Reset() + r.statusCode = statusCode + cl := r.header.Get("content-length") + te := r.header.Get("transfer-encoding") + if cl == "" || te != "" { + if te == "" { + r.header.Set("transfer-encoding", "chunked") + } + r.w = &chunkWriter{r.buf} + } else { + r.w = r.buf + } + fmt.Fprintf(r.buf, "HTTP/1.1 %d %s\r\n", statusCode, + http.StatusText(statusCode)) + for k, v := range r.header { + fmt.Fprintf(r.buf, "%s: %s\r\n", strings.Title(k), strings.Join(v, ",")) + } + fmt.Fprintf(r.buf, "\r\n") + r.responseSend = true +} + +func (r *responseHandler) Header() http.Header { + return r.header +} + +func (r *responseHandler) Write(buf []byte) (int, error) { + if !r.responseSend { + r.WriteHeader(http.StatusOK) + } + n, err := r.w.Write(buf) + if r.buf.Len() > 2048 { + r.buf.WriteTo(r.c) + } + return n, err +} + +var _ http.ResponseWriter = &responseHandler{} + +type chunkWriter struct { + w io.Writer +} + +func (cw *chunkWriter) Write(buf []byte) (int, error) { + n := len(buf) + if n == 0 { + return fmt.Fprintf(cw.w, "0\r\n\r\n") + } + return fmt.Fprintf(cw.w, "%x\r\n%s\r\n", n, string(buf)) +} diff --git a/routers.go b/routers.go index 63e9c03..504f60b 100644 --- a/routers.go +++ b/routers.go @@ -1,8 +1,6 @@ package main import ( - "bufio" - "bytes" "fmt" "io" "log" @@ -18,7 +16,6 @@ import ( "github.com/fangdingjun/gnutls" auth "github.com/fangdingjun/go-http-auth" "github.com/fangdingjun/gofast" - nghttp2 "github.com/fangdingjun/nghttp2-go" loghandler "github.com/gorilla/handlers" "github.com/gorilla/mux" ) @@ -270,103 +267,3 @@ func (m myURLMatch) match(r *http.Request, route *mux.RouteMatch) bool { ret := m.re.MatchString(r.URL.Path) return ret } - -func handleHTTPClient(c net.Conn, handler http.Handler) { - tlsconn := c.(*gnutls.Conn) - if err := tlsconn.Handshake(); err != nil { - log.Println(err) - return - } - state := tlsconn.ConnectionState() - if state.NegotiatedProtocol == "h2" { - h2conn, err := nghttp2.Server(tlsconn, handler) - if err != nil { - log.Println(err) - } - h2conn.Run() - return - } - - defer c.Close() - r := bufio.NewReader(tlsconn) - buf := new(bytes.Buffer) - for { - req, err := http.ReadRequest(r) - if err != nil { - return - } - addr := tlsconn.RemoteAddr().String() - req.RemoteAddr = addr - rh := &responseHandler{ - c: tlsconn, - header: http.Header{}, - buf: buf, - } - handler.ServeHTTP(rh, req) - rh.Write(nil) - rh.buf.WriteTo(rh.c) - } -} - -type responseHandler struct { - c net.Conn - statusCode int - header http.Header - responseSend bool - w io.Writer - buf *bytes.Buffer -} - -func (r *responseHandler) WriteHeader(statusCode int) { - if r.responseSend { - return - } - r.buf.Reset() - r.statusCode = statusCode - cl := r.header.Get("content-length") - te := r.header.Get("transfer-encoding") - if cl == "" || te != "" { - if te == "" { - r.header.Set("transfer-encoding", "chunked") - } - r.w = &chunkWriter{r.buf} - } else { - r.w = r.buf - } - fmt.Fprintf(r.buf, "HTTP/1.1 %d %s\r\n", statusCode, - http.StatusText(statusCode)) - for k, v := range r.header { - fmt.Fprintf(r.buf, "%s: %s\r\n", strings.Title(k), strings.Join(v, ",")) - } - fmt.Fprintf(r.buf, "\r\n") - r.responseSend = true -} - -func (r *responseHandler) Header() http.Header { - return r.header -} - -func (r *responseHandler) Write(buf []byte) (int, error) { - if !r.responseSend { - r.WriteHeader(http.StatusOK) - } - n, err := r.w.Write(buf) - if r.buf.Len() > 2048 { - r.buf.WriteTo(r.c) - } - return n, err -} - -var _ http.ResponseWriter = &responseHandler{} - -type chunkWriter struct { - w io.Writer -} - -func (cw *chunkWriter) Write(buf []byte) (int, error) { - n := len(buf) - if n == 0 { - return fmt.Fprintf(cw.w, "0\r\n\r\n") - } - return fmt.Fprintf(cw.w, "%x\r\n%s\r\n", n, string(buf)) -}