out/cmd/testapi/main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/coder/websocket"
	"golang.org/x/time/rate"
)

func main() {
	ctx, can := signal.NotifyContext(context.Background(), syscall.SIGINT)
	defer can()

	if err := run(ctx); err != nil {
		panic(err)
	}
}

func run(ctx context.Context) error {
	fs := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
	port := fs.Int("p", 8080, "port")
	if err := fs.Parse(os.Args[1:]); err != nil {
		return err
	}

	S := &S{
		ctx:     ctx,
		limiter: rate.NewLimiter(10, 10),
	}

	s := &http.Server{
		Addr:    fmt.Sprintf(":%d", *port),
		Handler: S,
	}

	go func() {
		<-ctx.Done()
		ctx, can := context.WithTimeout(context.Background(), time.Second)
		defer can()
		s.Shutdown(ctx)
	}()

	log.Println("listening on", *port)
	if err := s.ListenAndServe(); err != nil && ctx.Err() == nil {
		return err
	}
	log.Println("shut down")
	return nil
}

type S struct {
	ctx     context.Context
	limiter *rate.Limiter
}

func (s *S) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if err := s.serveHTTP(w, r); err != nil {
		log.Println(r.URL.Path, "//", err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func (s *S) serveHTTP(w http.ResponseWriter, r *http.Request) error {
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Headers", "X-Auth-Token, content-type, Content-Type")
	if r.Method == http.MethodOptions {
		w.Header().Set("Content-Length", "0")
		w.Header().Set("Content-Type", "text/plain")
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, OPTIONS, TRACE, PATCH, HEAD, DELETE")
		w.WriteHeader(http.StatusOK)
		return nil
	}
	if isV1(r) || isWS(r) {
		return s.serveAPI(w, r)
	}
	return s.serveStatic(w, r)
}

func isV1(r *http.Request) bool {
	return strings.HasPrefix(r.URL.Path, "/v1/")
}

func isWS(r *http.Request) bool {
	return r.URL.Path == "/ws" || strings.HasPrefix(r.URL.Path, "/ws/")
}

func (s *S) serveStatic(w http.ResponseWriter, r *http.Request) error {
	return fmt.Errorf("not impl static")
}

func (s *S) serveAPI(w http.ResponseWriter, r *http.Request) error {
	if err := s.injectContext(w, r); err == io.EOF {
		http.Redirect(w, r, "/", http.StatusSeeOther)
		return nil
	} else if err != nil {
		return err
	}

	if isWS(r) {
		return s.serveWS(w, r)
	} else if isV1(r) {
		return s.serveV1(w, r)
	}

	http.NotFound(w, r)
	return nil
}

type Session struct {
	ID string
}

func (s *S) injectContext(w http.ResponseWriter, r *http.Request) error {
	id := r.URL.Query().Get("uuid")
	if id == "" {
		c, err := r.Cookie("uuid")
		if err != nil || c.Value == "" {
			return io.EOF
		}
		id = c.Value
	}
	ctx := r.Context()
	ctx = context.WithValue(ctx, "session", Session{
		ID: id,
	})
	*r = *r.WithContext(ctx)
	return nil
}

func (s *S) Session(ctx context.Context) Session {
	v, _ := ctx.Value("session").(Session)
	return v
}

func (s *S) serveWS(w http.ResponseWriter, r *http.Request) error {
	ctx := r.Context()

	c, err := websocket.Accept(w, r, nil)
	if err != nil {
		return err
	}
	defer c.CloseNow()

	for {
		select {
		case <-ctx.Done():
			return nil
		case <-time.After(time.Second):
			switch r.URL.Path {
			case "/ws/page/a":
				if err := c.Write(ctx, 1, []byte(`{"page":"a", "items": [{"name":"x"}, {"name":"y"}]}`)); err != nil {
					return err
				}
			case "/ws/page/a/event/a":
				if err := c.Write(ctx, 1, []byte(`{"page":"a", "event":"a", "items": [{"name":"x"}, {"name":"y"}]}`)); err != nil {
					return err
				}
			case "/ws/page/b/event/a":
				if err := c.Write(ctx, 1, []byte(`{"page":"b", "event":"a", "items": [{"name":"x", "title":"x1", "tags":[{"k":"k", "v":"v"}, {"k":"K", "v":"V"}, {"k":"kkk","v":"vvv"}]}, {"name":"y", "tags":[{"k":"k", "v":"v"}, {"k":"K", "v":"V"}, {"k":"kkk","v":"vvv"}]}], "help": "hello world"}`)); err != nil {
					return err
				}
			case "/ws/page/b/event/b":
				if err := c.Write(ctx, 1, []byte(`{"page":"b", "event":"b", "items": [{"name":"x", "title":"x1", "tags":[{"k":"k", "v":"v"}, {"k":"K", "v":"V"}, {"k":"kkk","v":"vvv"}]}, {"name":"y", "tags":[{"k":"k", "v":"v"}, {"k":"K", "v":"V"}, {"k":"kkk","v":"vvv"}]}], "help": "hello world"}`)); err != nil {
					return err
				}
			default:
				if err := c.Write(ctx, 1, []byte(`hello world`)); err != nil {
					return err
				}
			}
		}
	}

	return nil
}

func (s *S) serveV1(w http.ResponseWriter, r *http.Request) error {
	ctx := r.Context()
	switch r.Method + r.URL.Path {
	case "GET/v1/state/" + s.Session(ctx).ID:
		if rand.Int()%2 == 0 {
			w.Write([]byte(`{"name": "foo"}`))
		} else {
			w.Write([]byte(`{"name": "bar", "party": "party name"}`))
		}
	case "PUT/v1/state/" + s.Session(ctx).ID + "/party":
		w.Write([]byte(`{}`))
	default:
		http.NotFound(w, r)
		return nil
	}
	return nil
}