Store oauth session for redirect, invite un-authed

Also added Dockerfile :)

Dockerfile

@@ -0,0 +1,23 @@
+FROM golang:alpine as builder
+
+WORKDIR /build/wlm
+COPY go.mod go.sum ./
+RUN go mod download
+
+COPY . .
+
+RUN apk add --update gcc musl-dev
+WORKDIR /build/wlm
+RUN go build -o wlm -ldflags "-s -w"
+
+
+FROM alpine
+
+WORKDIR /app
+EXPOSE 8080
+
+RUN apk add --no-cache ca-certificates sqlite
+
+COPY --from=builder /build/wlm/wlm wlm
+
+ENTRYPOINT ["/app/wlm"]

auth/auth.go

@@ -49,14 +49,14 @@ type MinecraftAuth struct {
 	IdentityToken string `json:"identityToken"`
 }
 
-func LoginUrl() string {
+func LoginUrl(stateId string) string {
 	oaConfig := &oauth2.Config{
 		ClientID:     os.Getenv("AZURE_OAUTH_CLIENT_ID"),
 		ClientSecret: os.Getenv("AZURE_OAUTH_CLIENT_SECRET"),
 		Endpoint:     endpoints.Microsoft,
 		Scopes:       []string{"Xboxlive.signin", "Xboxlive.offline_access"},
 	}
-	return oaConfig.AuthCodeURL("foo")
+	return oaConfig.AuthCodeURL(stateId)
 }
 
 func Authenticate(code string) (store.User, error) {

main.go

@@ -24,11 +24,12 @@ func main() {
 	// API endpoints
 	mux.Group(func(mux *flow.Mux) {
 		mux.Use(handlers.Cors)
+		mux.HandleFunc("/api/v1/invite/:id", handlers.GetInvite, "GET", "DELETE")
+
 		mux.Use(handlers.SessionAuth)
 		mux.HandleFunc("/api/v1/me", handlers.CurrentUser, "GET")
 
 		mux.HandleFunc("/api/v1/invites", handlers.CreateInvite, "POST")
-		mux.HandleFunc("/api/v1/invite/:id", handlers.GetInvite, "GET", "DELETE")
 		mux.HandleFunc("/api/v1/invite/:id/accept", handlers.AcceptInvite, "POST")
 
 		mux.HandleFunc("/api/v1/servers", handlers.Server, "GET", "POST")

store/database.go

@@ -32,6 +32,14 @@ type Storer interface {
 	SaveUser(user User) error
 	SaveSession(token string, user User) error
 	SessionUser(token string) (User, error)
+	SaveOauthState(state OauthState) error
+	OauthState(id string) (OauthState, error)
+}
+
+type OauthState struct {
+	Id     string
+	Origin string
+	State  string
 }
 
 type Invite struct {
@@ -301,3 +309,43 @@ func (s *Store) SessionUser(token string) (User, error) {
 
 	return s.GetUser(sess.UID)
 }
+
+func (s *Store) SaveOauthState(state OauthState) error {
+	_, err := s.OauthState(state.Id)
+	if err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			q, err := s.database.Prepare("INSERT INTO oauth_states (state_id, origin, state) VALUES ($1, $2, $3)")
+			if err != nil {
+				return err
+			}
+			_, err = q.Exec(state.Id, state.Origin, state.State)
+
+			if err != nil {
+				return err
+			}
+			return nil
+		} else {
+			return err
+		}
+	}
+	q, err := s.database.Prepare("UPDATE oauth_states SET state=$2 WHERE state_id=$1")
+	if err != nil {
+		return err
+	}
+	_, err = q.Exec(state.Id, state.State)
+	if err != nil {
+		return err
+	}
+
+	return err
+}
+
+func (s *Store) OauthState(id string) (OauthState, error) {
+	q := s.database.QueryRow("SELECT * FROM oauth_states WHERE state_id=$1", id)
+	var state OauthState
+	err := q.Scan(&state.Id, &state.Origin, &state.State)
+	if err != nil {
+		return OauthState{}, err
+	}
+	return state, nil
+}

store/database.sql

@@ -44,4 +44,11 @@ CREATE TABLE `invite_log` (
     `user` TEXT NOT NULL,
     PRIMARY KEY (`entry_id`),
     FOREIGN KEY (`invite`) REFERENCES invites(`token`)
-)
+) STRICT;
+
+CREATE TABLE `oauth_states` (
+    `state_id` TEXT NOT NULL,
+    `origin` TEXT NOT NULL,
+    `state` TEXT NOT NULL,
+    PRIMARY KEY (`state_id`)
+) STRICT;

transport/http.go

@@ -170,11 +170,28 @@ func (h *Handler) AcceptInvite(w http.ResponseWriter, r *http.Request) {
 }
 
 func (h *Handler) AuthRedirect(w http.ResponseWriter, r *http.Request) {
-	http.Redirect(w, r, auth.LoginUrl(), http.StatusTemporaryRedirect)
+	state := store.OauthState{
+		Id:     uuid.New().String(),
+		Origin: r.Header.Get("Referer"),
+		State:  "started",
+	}
+	err := h.store.SaveOauthState(state)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+	http.Redirect(w, r, auth.LoginUrl(state.Id), http.StatusTemporaryRedirect)
 }
 
 func (h *Handler) AuthCallback(w http.ResponseWriter, r *http.Request) {
 	code := r.URL.Query()["code"][0]
+	stateId := r.URL.Query()["state"][0]
+	state, err := h.store.OauthState(stateId)
+	if err != nil || state.State == "completed" {
+		http.Error(w, err.Error(), http.StatusForbidden)
+		return
+	}
+
 	user, err := auth.Authenticate(code)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
@@ -202,8 +219,14 @@ func (h *Handler) AuthCallback(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
 	}
+	state.State = "completed"
+	err = h.store.SaveOauthState(state)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
 
-	http.Redirect(w, r, "/", http.StatusTemporaryRedirect)
+	http.Redirect(w, r, state.Origin, http.StatusTemporaryRedirect)
 }
 
 func generateSessionToken() (string, error) {