package main import ( "context" "flag" "fmt" "io" "log" "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 { 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, err := r.Cookie("uuid") if err != nil || id.Value == "" { return io.EOF } ctx := r.Context() ctx = context.WithValue(ctx, "session", Session{ ID: id.Value, }) *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 "PUT/v1/state/" + s.Session(ctx).ID + "/party": return nil } http.NotFound(w, r) return nil }