diff --git a/unix/server.go b/unix/server.go index c1fa085144cdd1269ebb04ea84c9c20b92f1d19e..19fd813900e22eaa541a181d529c5b88bc03b05a 100644 --- a/unix/server.go +++ b/unix/server.go @@ -1,6 +1,7 @@ package unix import ( + "bufio" "context" "errors" "io" @@ -21,11 +22,6 @@ type Handler interface { ServeConnection(c net.Conn) } -// HandlerFunc is a function adapter for Handler. -type HandlerFunc func(net.Conn) - -func (f HandlerFunc) ServeConnection(c net.Conn) { f(c) } - // SocketServer accepts connections on a UNIX socket, speaking the // line-based wire protocol, and dispatches incoming requests to the // wrapped Server. @@ -134,7 +130,22 @@ func (s *SocketServer) Serve() error { // LineHandler is the handler for LineServer. type LineHandler interface { - ServeLine(context.Context, []byte) (string, error) + ServeLine(context.Context, LineResponseWriter, []byte) error +} + +// ErrCloseConnection must be returned by a LineHandler when we want +// to cleanly terminate the connection without raising an error. +var ErrCloseConnection = errors.New("close") + +// LineResponseWriter writes a single-line response to the underlying +// connection. +type LineResponseWriter interface { + // WriteLine writes a response (which must include the + // line terminator). + WriteLine([]byte) error + + // WriteLineCRLF writes a response and adds a line terminator. + WriteLineCRLF([]byte) error } // LineServer implements a line-based text protocol. It satisfies the @@ -164,8 +175,32 @@ func NewLineServer(h LineHandler) *LineServer { } } +var crlf = []byte{'\r', '\n'} + +type lrWriter struct { + *bufio.Writer +} + +func (w *lrWriter) WriteLine(data []byte) error { + if _, err := w.Writer.Write(data); err != nil { + return err + } + return w.Writer.Flush() +} + +func (w *lrWriter) WriteLineCRLF(data []byte) error { + if _, err := w.Writer.Write(data); err != nil { + return err + } + if _, err := w.Writer.Write(crlf); err != nil { + return err + } + return w.Writer.Flush() +} + func (l *LineServer) ServeConnection(nc net.Conn) { c := textproto.NewConn(nc) + rw := &lrWriter{bufio.NewWriter(nc)} for { nc.SetReadDeadline(time.Now().Add(l.IdleTimeout)) line, err := c.ReadLineBytes() @@ -176,21 +211,20 @@ func (l *LineServer) ServeConnection(nc net.Conn) { break } - // Create a context for the request and call the handler with it. + // Create a context for the request and call the + // handler with it. Set a write deadline on the + // connection to allow the full RequestTimeout time to + // generate the response. + nc.SetWriteDeadline(time.Now().Add(l.RequestTimeout + l.WriteTimeout)) ctx, cancel := context.WithTimeout(context.Background(), l.RequestTimeout) - response, err := l.handler.ServeLine(ctx, line) + err = l.handler.ServeLine(ctx, rw, line) cancel() // Close the connection on error, or on empty response. - if response != "" { - nc.SetWriteDeadline(time.Now().Add(l.WriteTimeout)) - c.PrintfLine(response) - } if err != nil { - log.Printf("request error: %v", err) - break - } - if response == "" { + if err != ErrCloseConnection { + log.Printf("request error: %v", err) + } break } }