Simplify ResponseWriter implementation

3 files changed, 91 insertions, 110 deletions

diff --git a/fs.go b/fs.go
index 6afb917..a7cbd57 100644
--- a/fs.go
+++ b/fs.go
@@ -39,7 +39,7 @@ func (fsh fsHandler) Respond(w *ResponseWriter, r *Request) {
 	// Detect mimetype
 	ext := path.Ext(p)
 	mimetype := mime.TypeByExtension(ext)
-	w.SetMediaType(mimetype)
+	w.Meta(mimetype)
 	// Copy file to response writer
 	_, _ = io.Copy(w, f)
 }
@@ -78,7 +78,7 @@ func ServeFile(w *ResponseWriter, fs FS, name string) {
 	// Detect mimetype
 	ext := path.Ext(name)
 	mimetype := mime.TypeByExtension(ext)
-	w.SetMediaType(mimetype)
+	w.Meta(mimetype)
 	// Copy file to response writer
 	_, _ = io.Copy(w, f)
 }
diff --git a/response.go b/response.go
index 66b5e6f..e603e3b 100644
--- a/response.go
+++ b/response.go
@@ -114,3 +114,92 @@ func (b *readCloserBody) Read(p []byte) (n int, err error) {
 	}
 	return b.ReadCloser.Read(p)
 }
+
+// ResponseWriter is used by a Gemini handler to construct a Gemini response.
+type ResponseWriter struct {
+	b           *bufio.Writer
+	status      Status
+	meta        string
+	setHeader   bool
+	wroteHeader bool
+	bodyAllowed bool
+}
+
+// NewResponseWriter returns a ResponseWriter that uses the provided io.Writer.
+func NewResponseWriter(w io.Writer) *ResponseWriter {
+	return &ResponseWriter{
+		b: bufio.NewWriter(w),
+	}
+}
+
+// Header sets the response header.
+func (w *ResponseWriter) Header(status Status, meta string) {
+	w.status = status
+	w.meta = meta
+}
+
+// Status sets the response status code.
+// It also sets the response meta to status.Meta().
+func (w *ResponseWriter) Status(status Status) {
+	w.status = status
+	w.meta = status.Meta()
+}
+
+// Meta sets the response meta.
+//
+// For successful responses, meta should contain the media type of the response.
+// For failure responses, meta should contain a short description of the failure.
+// The response meta should not be greater than 1024 bytes.
+func (w *ResponseWriter) Meta(meta string) {
+	w.meta = meta
+}
+
+// Write writes data to the connection as part of the response body.
+// If the response status does not allow for a response body, Write returns
+// ErrBodyNotAllowed.
+//
+// Write writes the response header if it has not already been written.
+// It writes a successful status code if one is not set.
+func (w *ResponseWriter) Write(b []byte) (int, error) {
+	if !w.wroteHeader {
+		w.writeHeader(StatusSuccess)
+	}
+	if !w.bodyAllowed {
+		return 0, ErrBodyNotAllowed
+	}
+	return w.b.Write(b)
+}
+
+func (w *ResponseWriter) writeHeader(defaultStatus Status) {
+	status := w.status
+	if status == 0 {
+		status = defaultStatus
+	}
+
+	meta := w.meta
+	if status.Class() == StatusClassSuccess {
+		w.bodyAllowed = true
+
+		if meta == "" {
+			meta = "text/gemini"
+		}
+	}
+
+	w.b.WriteString(strconv.Itoa(int(status)))
+	w.b.WriteByte(' ')
+	w.b.WriteString(meta)
+	w.b.Write(crlf)
+	w.wroteHeader = true
+}
+
+// Flush writes any buffered data to the underlying io.Writer.
+//
+// Flush writes the response header if it has not already been written.
+// It writes a failure status code if one is not set.
+func (w *ResponseWriter) Flush() error {
+	if !w.wroteHeader {
+		w.writeHeader(StatusTemporaryFailure)
+	}
+	// Write errors from writeHeader will be returned here.
+	return w.b.Flush()
+}
diff --git a/server.go b/server.go
index ecfd62d..8223cbf 100644
--- a/server.go
+++ b/server.go
@@ -1,11 +1,8 @@
 package gemini
 
 import (
-	"bufio"
 	"crypto/tls"
 	"errors"
-	"fmt"
-	"io"
 	"log"
 	"net"
 	"strings"
@@ -236,111 +233,6 @@ func (s *Server) logf(format string, args ...interface{}) {
 	}
 }
 
-// ResponseWriter is used by a Gemini handler to construct a Gemini response.
-type ResponseWriter struct {
-	status      Status
-	meta        string
-	b           *bufio.Writer
-	bodyAllowed bool
-	wroteHeader bool
-	mediatype   string
-}
-
-// NewResponseWriter returns a ResponseWriter that uses the provided io.Writer.
-func NewResponseWriter(w io.Writer) *ResponseWriter {
-	return &ResponseWriter{
-		b: bufio.NewWriter(w),
-	}
-}
-
-// Header sets the response header.
-//
-// Meta contains more information related to the response status.
-// For successful responses, Meta should contain the mimetype of the response.
-// For failure responses, Meta should contain a short description of the failure.
-// Meta should not be longer than 1024 bytes.
-func (w *ResponseWriter) Header(status Status, meta string) {
-	w.status = status
-	w.meta = meta
-}
-
-// Status sets the response header to the given status code.
-//
-// Status is equivalent to Header(status, status.Meta())
-func (w *ResponseWriter) Status(status Status) {
-	w.status = status
-	w.meta = status.Meta()
-}
-
-// SetMediaType sets the media type that will be written for a successful response.
-// If the mimetype is not set, it will default to "text/gemini".
-func (w *ResponseWriter) SetMediaType(mediatype string) {
-	w.mediatype = mediatype
-}
-
-// Write writes data to the connection as part of the response body.
-// If the response status does not allow for a response body, Write returns
-// ErrBodyNotAllowed.
-//
-// If the response header has not yet been written, Write calls WriteHeader
-// with StatusSuccess and the mimetype set in SetMimetype.
-func (w *ResponseWriter) Write(b []byte) (int, error) {
-	if !w.wroteHeader {
-		err := w.writeHeader()
-		if err != nil {
-			return 0, err
-		}
-	}
-
-	if !w.bodyAllowed {
-		return 0, ErrBodyNotAllowed
-	}
-
-	return w.b.Write(b)
-}
-
-func (w *ResponseWriter) writeHeader() error {
-	status := w.status
-	if status == 0 {
-		status = StatusSuccess
-	}
-
-	meta := w.meta
-
-	if status.Class() == StatusClassSuccess {
-		w.bodyAllowed = true
-
-		if meta == "" {
-			meta = w.mediatype
-		}
-
-		if meta == "" {
-			meta = "text/gemini"
-		}
-	}
-
-	_, err := fmt.Fprintf(w.b, "%d %s\r\n", int(status), meta)
-	if err != nil {
-		return fmt.Errorf("failed to write response header: %w", err)
-	}
-
-	w.wroteHeader = true
-
-	return nil
-}
-
-// Flush writes any buffered data to the underlying io.Writer.
-func (w *ResponseWriter) Flush() error {
-	if !w.wroteHeader {
-		err := w.writeHeader()
-		if err != nil {
-			return err
-		}
-	}
-
-	return w.b.Flush()
-}
-
 // A Responder responds to a Gemini request.
 type Responder interface {
 	// Respond accepts a Request and constructs a Response.