💾 Archived View for source.community › ckaznocha › gemini › raw › main › response.go captured on 2024-05-10 at 11:19:36.

View Raw

More Information

⬅️ Previous capture (2021-12-17)

-=-=-=-=-=-=-

package gemini

import (
	"context"
	"io"
	"mime"
)

// ResponseWriter is interface to interact with a Gemini response. Calling any
// of its meethods after one has already been called or after
// Handler.ServeGemini has returned is a no-op.
type ResponseWriter interface {
	Failure(ctx context.Context, code StatusCode, msg string)
	Input(ctx context.Context, prompt string, isSensitive bool)
	Redirect(ctx context.Context, redirectURL string, isPermanant bool)
	Success(ctx context.Context, mimeType string) io.Writer
}

type response struct {
	conn io.Writer
	sent atomicBool
}

//nolint:gochecknoglobals // allocate some frequently used byte slices.
var (
	responseHeaderSpace          = []byte{' '}
	responseHeaderLineTerminator = []byte("\r\n")
)

func (r *response) writeHeader(ctx context.Context, code StatusCode, meta []byte) {
	if r.sent.isSet() {
		return
	}
	defer r.sent.set()

	logWriter := func(string, bool) {}

	srv, ok := ServerFromCtx(ctx)
	if ok && srv.LogHandler != nil {
		logWriter = srv.LogHandler
	}

	c, err := code.MarshalText()
	if err != nil {
		logWriter(err.Error(), true)

		return
	}

	if _, err = r.conn.Write(c); err != nil {
		logWriter(err.Error(), true)

		return
	}

	if _, err = r.conn.Write(responseHeaderSpace); err != nil {
		logWriter(err.Error(), true)

		return
	}

	if _, err = r.conn.Write(meta); err != nil {
		logWriter(err.Error(), true)

		return
	}

	if _, err = r.conn.Write(responseHeaderLineTerminator); err != nil {
		logWriter(err.Error(), true)

		return
	}
}

func (r *response) Input(ctx context.Context, prompt string, isSensitive bool) {
	code := StatusInput
	if isSensitive {
		code = StatusSensitiveInput
	}

	r.writeHeader(ctx, code, []byte(prompt))
}

func (r *response) Redirect(ctx context.Context, redirectURL string, isPermanant bool) {
	code := StatusTemporaryRedirect
	if isPermanant {
		code = StatusPermanentRedirect
	}

	r.writeHeader(ctx, code, []byte(redirectURL))
}

func (r *response) Success(ctx context.Context, mimeType string) io.Writer {
	if mimeType == "" {
		mimeType = mime.FormatMediaType("text/gemini", map[string]string{"charset": "utf-8", "lang": "en"})
	}

	r.writeHeader(ctx, StatusSuccess, []byte(mimeType))

	return r.conn
}

func (r *response) Failure(ctx context.Context, code StatusCode, msg string) {
	category := code.ToCategory()
	if category < StatusCategoryTemporaryFailure || category > StatusCategoryPermanentFailure {
		code = StatusPermanentFailure
	}

	r.writeHeader(ctx, code, []byte(msg))
}