Add SetMimetype convenience function

1 files changed, 26 insertions, 1 deletions

diff --git a/server.go b/server.go
index 8187e71..38e7ce8 100644
--- a/server.go
+++ b/server.go
@@ -95,6 +95,8 @@ func (s *Server) Serve(l net.Listener) error {
 type ResponseWriter struct {
 	w           *bufio.Writer
 	bodyAllowed bool
+	wroteHeader bool
+	mimetype    string
 }
 
 func newResponseWriter(conn net.Conn) *ResponseWriter {
@@ -104,12 +106,16 @@ func newResponseWriter(conn net.Conn) *ResponseWriter {
 }
 
 // WriteHeader writes the response header.
+// If the header has already been written, WriteHeader does nothing.
 //
 // 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 (r *ResponseWriter) WriteHeader(status int, meta string) {
+	if r.wroteHeader {
+		return
+	}
 	r.w.WriteString(strconv.Itoa(status))
 	r.w.WriteByte(' ')
 	r.w.WriteString(meta)
@@ -119,13 +125,32 @@ func (r *ResponseWriter) WriteHeader(status int, meta string) {
 	if status/10 == StatusClassSuccess {
 		r.bodyAllowed = true
 	}
+	r.wroteHeader = true
+}
+
+// SetMimetype sets the mimetype that will be written for a successful response.
+// The provided mimetype will only be used if Write is called without calling
+// WriteHeader.
+// If the mimetype is not set, it will default to "text/gemini".
+func (r *ResponseWriter) SetMimetype(mimetype string) {
+	r.mimetype = mimetype
 }
 
 // Write writes the response body.
 // If the response status does not allow for a response body, Write returns
 // ErrBodyNotAllowed.
-// WriteHeader must be called before Write.
+//
+// If WriteHeader has not yet been called, Write calls
+// WriteHeader(StatusSuccess, mimetype) where mimetype is the mimetype set in
+// SetMimetype. If no mimetype is set, a default of "text/gemini" will be used.
 func (r *ResponseWriter) Write(b []byte) (int, error) {
+	if !r.wroteHeader {
+		if r.mimetype != "" {
+			r.WriteHeader(StatusSuccess, r.mimetype)
+		} else {
+			r.WriteHeader(StatusSuccess, "text/gemini")
+		}
+	}
 	if !r.bodyAllowed {
 		return 0, ErrBodyNotAllowed
 	}