~adnano/astronaut

ref: 72ba913688a9d2dafb0bfdf11a5cee16685ce98a astronaut/response.go -rw-r--r-- 1.4 KiB
72ba9136Adnan Maolood Use scfg to configure settings 4 months ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
package main

import (
	"bytes"
	"io/ioutil"

	"git.sr.ht/~adnano/go-gemini"
)

// ResponseWriter implements a gemini.ResponseWriter that constructs
// a Gemini response.
type ResponseWriter struct {
	status      gemini.Status
	meta        string
	body        bytes.Buffer
	mediatype   string
	wroteHeader bool
	bodyAllowed bool
}

// SetMediaType implements the gemini.ResponseWriter interface.
func (w *ResponseWriter) SetMediaType(mediatype string) {
	w.mediatype = mediatype
}

// Write implements the gemini.ResponseWriter interface.
func (w *ResponseWriter) Write(b []byte) (int, error) {
	if !w.wroteHeader {
		meta := w.mediatype
		if meta == "" {
			// Use default media type
			meta = "text/gemini; charset=utf-8"
		}
		w.WriteHeader(gemini.StatusSuccess, meta)
	}
	if !w.bodyAllowed {
		return 0, gemini.ErrBodyNotAllowed
	}
	return w.body.Write(b)
}

// WriteHeader implements the gemini.ResponseWriter interface.
func (w *ResponseWriter) WriteHeader(status gemini.Status, meta string) {
	if w.wroteHeader {
		return
	}

	if status.Class() == gemini.StatusSuccess {
		w.bodyAllowed = true
	}

	w.status = status
	w.meta = meta
	w.wroteHeader = true
}

// Response returns the constructed response.
func (w *ResponseWriter) Response() *gemini.Response {
	return &gemini.Response{
		Status: w.status,
		Meta:   w.meta,
		Body:   ioutil.NopCloser(&w.body),
	}
}

func (w *ResponseWriter) Flush() error {
	return nil
}