wpw-final/internal/http/server.go

package http

import (
	"bytes"
	"context"
	"mime"
	"net/http"
	"strings"
	"time"

	"wpw-common/internal"
	"wpw-common/public"

	"github.com/apache/thrift/lib/go/thrift"
	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/fiber/v2/middleware/cors"
	"github.com/gofiber/fiber/v2/middleware/filesystem"
	"github.com/gofiber/fiber/v2/middleware/recover"
	"github.com/gofiber/fiber/v2/middleware/session"
	"go.uber.org/zap"
)

var (
	defaultServer = newServer()
	Session       = session.New(session.Config{
		KeyLookup: "header:X-Hub-Track-Token",
	})

	thriftContentType = []byte("application/x-thrift")

	contextKeyFiber contextKey = "fiber-app"

	hubHeader = []string{
		"X-Hub-Env",
		"X-Hub-App",
		"X-Hub-Version",
		"X-Hub-App",
		"X-Hub-Track-Token",
		"X-Hub-Auth-Token",
	}
)

type contextKey string

type server struct {
	pCtx   context.Context
	app    *fiber.App
	logger *zap.SugaredLogger
}

func newServer() *server {
	app := fiber.New()

	app.Use(recover.New())

	app.Use(cors.New(cors.Config{
		// AllowOrigins: "https://gofiber.io, https://gofiber.net",
		AllowOrigins: "*",
		AllowHeaders: strings.Join(append([]string{
			"Origin",
			"Content-Type",
			"Accept",
		}, hubHeader...), ", "),
		ExposeHeaders: strings.Join(hubHeader, ", "),
	}))

	app.Use(filesystem.New(filesystem.Config{
		Root:         http.FS(public.FS),
		Browse:       false,
		Index:        "index.html",
		MaxAge:       3000,
		NotFoundFile: "index.html",
	}))

	srv := &server{app: app}
	return srv
}

func (s *server) registerService(path string, p thrift.TProcessor) {
	processorFactory := thrift.NewTProcessorFactory(p)
	protocolFactory := thrift.NewTCompactProtocolFactoryConf(&thrift.TConfiguration{})
	contentTypeResp := mime.FormatMediaType(string(thriftContentType), map[string]string{
		"charset":  "utf-8",
		"protocol": "TCOMPACT",
	})

	s.app.Post(path, func(c *fiber.Ctx) error {
		mBegin := time.Now()
		var mDur time.Duration

		c.Response().Header.Set("Content-Type", contentTypeResp)
		if !bytes.HasPrefix(c.Request().Header.ContentType(), thriftContentType) {
			c.Response().Header.Set("x-error-type", "bad-request")
			c.Response().Header.Set("x-error-message", "bad request")
			return c.SendStatus(400)
		}

		body := c.Context().Request.Body()
		transport := thrift.NewStreamTransport(bytes.NewReader(body), c)
		protocol := protocolFactory.GetProtocol(transport)

		// put fiber ctx to the context
		ctx := WithFiberContext(s.pCtx, c)

		// process method
		handled, err := processorFactory.GetProcessor(transport).Process(ctx, protocol, protocol)
		if err != nil {
			s.logger.Errorw("process method", "err", err)
			switch err.(type) {
			case thrift.TTransportException:
				c.Response().Header.Set("x-error-type", "transport-exception")
				c.Response().Header.Set("x-error-message", err.Error())
				goto internal_error
			case thrift.TProtocolException:
				c.Response().Header.Set("x-error-type", "protocol-exception")
				c.Response().Header.Set("x-error-message", err.Error())
				goto internal_error
			case thrift.TApplicationException:
				c.Response().Header.Set("x-error-type", "application-exception")
				c.Response().Header.Set("x-error-message", err.Error())
			default:
				c.Response().Header.Set("x-error-type", "exception")
				c.Response().Header.Set("x-error-message", err.Error())
			}
		}

		mDur = time.Since(mBegin)
		c.Response().Header.Set("x-trace-duration", mDur.String())

		_ = handled
		if err := protocol.Flush(c.Context()); err != nil {
			return err
		}
		return c.SendStatus(200)
	internal_error:
		// do something!
		return c.SendStatus(200)
	})
}

func (s *server) run(appCtx *internal.AppContext, addr string) error {
	s.pCtx = internal.WithAppContext(context.Background(), appCtx)
	s.logger = appCtx.Logger.Sugar().Named("httpsrv")
	s.logger.Infow("starting server",
		"addr", addr)
	return s.app.Listen(addr)
}

func Router() *fiber.App {
	return defaultServer.app
}

func GetFiberFromContext(ctx context.Context) *fiber.Ctx {
	return ctx.Value(contextKeyFiber).(*fiber.Ctx)
}

func WithFiberContext(ctx context.Context, c *fiber.Ctx) context.Context {
	return context.WithValue(ctx, contextKeyFiber, c)
}

func RegisterService(path string, p thrift.TProcessor) {
	defaultServer.registerService(path, p)
}

func Run(appCtx *internal.AppContext, addr string) error {
	return defaultServer.run(appCtx, addr)
}