Try websockets

api/api.go

@@ -3,168 +3,110 @@ package api
 import (
 	"encoding/json"
 	"github.com/1f349/lotus/imap"
-	"github.com/1f349/lotus/imap/marshal"
+	"github.com/gorilla/websocket"
-	imap2 "github.com/emersion/go-imap"
 	"github.com/julienschmidt/httprouter"
 	"log"
 	"net/http"
 	"time"
 )
 
-func SetupApiServer(listen string, auth func(callback AuthCallback) httprouter.Handle, send Smtp, recv Imap) *http.Server {
+var upgrader = websocket.Upgrader{}
+
+func SetupApiServer(listen string, auth *AuthChecker, send Smtp, recv Imap) *http.Server {
 	r := httprouter.New()
 
 	// === ACCOUNT ===
-	r.GET("/identities", auth(func(rw http.ResponseWriter, req *http.Request, params httprouter.Params, b AuthClaims) {
+	r.GET("/identities", auth.Middleware(func(rw http.ResponseWriter, req *http.Request, params httprouter.Params, b AuthClaims) {
 		// TODO(melon): find users aliases and other account data
 	}))
 
 	// === SMTP ===
-	r.POST("/message", auth(MessageSender(send)))
+	r.POST("/smtp", auth.Middleware(MessageSender(send)))
 
-	r.Handle(http.MethodConnect, "/", func(rw http.ResponseWriter, req *http.Request, params httprouter.Params) {
+	r.Handle(http.MethodConnect, "/imap", func(rw http.ResponseWriter, req *http.Request, params httprouter.Params) {
-		
+		// upgrade to websocket conn and defer close
+		c, err := upgrader.Upgrade(rw, req, nil)
+		if err != nil {
+			log.Println("[Imap] Failed to upgrade to websocket:", err)
+			return
+		}
+		defer c.Close()
+
+		// set a really short deadline to refuse unauthenticated clients
+		deadline := time.Now().Add(5 * time.Second)
+		_ = c.SetReadDeadline(deadline)
+		_ = c.SetWriteDeadline(deadline)
+
+		// close on all possible errors, assume we are being attacked
+		mt, msg, err := c.ReadMessage()
+		if err != nil {
+			return
+		}
+		if mt != websocket.TextMessage {
+			return
+		}
+		if len(msg) >= 2000 {
+			return
+		}
+
+		// get a "possible" auth token value
+		authToken := string(msg)
+
+		// wait for authToken or error
+		// exit on empty reply
+		if authToken == "" {
+			return
+		}
+
+		// check the token
+		authUser, err := auth.Check(authToken)
+		if err != nil {
+			// exit on error
+			return
+		}
+
+		_ = authUser
+
+		client, err := recv.MakeClient(authUser.Subject)
+		if err != nil {
+			_ = c.WriteJSON(map[string]string{"Error": "Making client failed"})
+			return
+		}
+
+		for {
+			// authenticated users get longer to reply
+			// a simple ping/pong setup bypasses this
+			d := time.Now().Add(5 * time.Minute)
+			_ = c.SetReadDeadline(d)
+			_ = c.SetWriteDeadline(d)
+
+			// read incoming message
+			var m struct {
+				Action string   `json:"action"`
+				Args   []string `json:"args"`
+			}
+			err := c.ReadJSON(&m)
+			if err != nil {
+				// errors should close the connection
+				return
+			}
+
+			// handle action
+			j, err := client.HandleWS(m.Action, m.Args)
+			if err != nil {
+				// errors should close the connection
+				return
+			}
+
+			// write outgoing message
+			err = c.WriteJSON(j)
+			if err != nil {
+				// errors should close the connection
+				return
+			}
+		}
 	})
 
-	// === IMAP ===
-	type mailboxStatusJson struct {
-		Folder string `json:"folder"`
-	}
-	r.GET("/mailbox/status", auth(imapClient(recv, func(rw http.ResponseWriter, req *http.Request, params httprouter.Params, cli *imap.Client, t mailboxStatusJson) error {
-		status, err := cli.Status(t.Folder)
-		if err != nil {
-			return err
-		}
-		return json.NewEncoder(rw).Encode(status)
-	})))
-
-	type mailboxListJson struct {
-		Folder  string `json:"folder"`
-		Pattern string `json:"pattern"`
-	}
-	r.GET("/mailbox/list", auth(imapClient(recv, func(rw http.ResponseWriter, req *http.Request, params httprouter.Params, cli *imap.Client, t mailboxListJson) error {
-		list, err := cli.List(t.Folder, t.Pattern)
-		if err != nil {
-			return err
-		}
-		return json.NewEncoder(rw).Encode(list)
-	})))
-
-	type mailboxCreateJson struct {
-		Name string `json:"name"`
-	}
-	r.POST("/mailbox/create", auth(imapClient(recv, func(rw http.ResponseWriter, req *http.Request, params httprouter.Params, cli *imap.Client, t mailboxCreateJson) error {
-		err := cli.Create(t.Name)
-		if err != nil {
-			return err
-		}
-		return json.NewEncoder(rw).Encode(map[string]string{"Status": "OK"})
-	})))
-
-	type messagesListJson struct {
-		Folder string `json:"folder"`
-		Start  uint32 `json:"start"`
-		End    uint32 `json:"end"`
-		Limit  uint32 `json:"limit"`
-	}
-	r.GET("/list-messages", auth(imapClient(recv, func(rw http.ResponseWriter, req *http.Request, params httprouter.Params, cli *imap.Client, t messagesListJson) error {
-		messages, err := cli.Fetch(t.Folder, t.Start, t.End, t.Limit)
-		if err != nil {
-			return err
-		}
-		return json.NewEncoder(rw).Encode(marshal.MessageSliceJson(messages))
-	})))
-
-	type messagesSearchJson struct {
-		Folder       string `json:"folder"`
-		SeqNum       imap2.SeqSet
-		Uid          imap2.SeqSet
-		Since        time.Time
-		Before       time.Time
-		SentSince    time.Time
-		SentBefore   time.Time
-		Body         []string
-		Text         []string
-		WithFlags    []string
-		WithoutFlags []string
-		Larger       uint32
-		Smaller      uint32
-	}
-	r.GET("/search-messages", auth(imapClient(recv, func(rw http.ResponseWriter, req *http.Request, params httprouter.Params, cli *imap.Client, t messagesSearchJson) error {
-		status, err := cli.Search(&imap2.SearchCriteria{
-			SeqNum:       t.SeqNum,
-			Uid:          t.Uid,
-			Since:        time.Time{},
-			Before:       time.Time{},
-			SentSince:    time.Time{},
-			SentBefore:   time.Time{},
-			Header:       nil,
-			Body:         nil,
-			Text:         nil,
-			WithFlags:    nil,
-			WithoutFlags: nil,
-			Larger:       0,
-			Smaller:      0,
-			Not:          nil,
-			Or:           nil,
-		})
-		if err != nil {
-			return err
-		}
-		return json.NewEncoder(rw).Encode(status)
-	})))
-	r.POST("/update-messages-flags", auth(imapClient(recv, func(rw http.ResponseWriter, req *http.Request, params httprouter.Params, cli *imap.Client, t mailboxStatusJson) {
-		status, err := cli.Status(t.Folder)
-		if err != nil {
-			rw.WriteHeader(http.StatusForbidden)
-			return
-		}
-		_ = json.NewEncoder(rw).Encode(status)
-	})))
-	r.GET("/list-messages", auth(imapClient(recv, func(rw http.ResponseWriter, req *http.Request, params httprouter.Params, cli *imap.Client, t statusJson) {
-		messages, err := cli.Fetch(t.Folder, 1, 100, 100)
-		if err != nil {
-			rw.WriteHeader(http.StatusForbidden)
-			return
-		}
-		err = json.NewEncoder(rw).Encode(marshal.ListMessagesJson(messages))
-		if err != nil {
-			log.Println("list-messages json encode error:", err)
-		}
-	})))
-	r.GET("/search-messages", auth(imapClient(recv, func(rw http.ResponseWriter, req *http.Request, params httprouter.Params, cli *imap.Client, t statusJson) {
-		status, err := cli.Status(t.Folder)
-		if err != nil {
-			rw.WriteHeader(http.StatusForbidden)
-			return
-		}
-		_ = json.NewEncoder(rw).Encode(status)
-	})))
-	r.POST("/create-message", auth(imapClient(recv, func(rw http.ResponseWriter, req *http.Request, params httprouter.Params, cli *imap.Client, t statusJson) {
-		status, err := cli.Status(t.Folder)
-		if err != nil {
-			rw.WriteHeader(http.StatusForbidden)
-			return
-		}
-		_ = json.NewEncoder(rw).Encode(status)
-	})))
-	r.POST("/update-messages-flags", auth(imapClient(recv, func(rw http.ResponseWriter, req *http.Request, params httprouter.Params, cli *imap.Client, t statusJson) {
-		status, err := cli.Status(t.Folder)
-		if err != nil {
-			rw.WriteHeader(http.StatusForbidden)
-			return
-		}
-		_ = json.NewEncoder(rw).Encode(status)
-	})))
-	r.POST("/copy-messages", auth(imapClient(recv, func(rw http.ResponseWriter, req *http.Request, params httprouter.Params, cli *imap.Client, t statusJson) {
-		status, err := cli.Status(t.Folder)
-		if err != nil {
-			rw.WriteHeader(http.StatusForbidden)
-			return
-		}
-		_ = json.NewEncoder(rw).Encode(status)
-	})))
-
 	return &http.Server{
 		Addr:              listen,
 		Handler:           r,

api/auth.go

@@ -2,6 +2,7 @@ package api
 
 import (
 	"crypto/subtle"
+	"errors"
 	"github.com/1f349/violet/utils"
 	"github.com/MrMelon54/mjwt"
 	"github.com/MrMelon54/mjwt/auth"
@@ -9,53 +10,68 @@ import (
 	"net/http"
 )
 
+var (
+	ErrInvalidToken    = errors.New("invalid token")
+	ErrInvalidAudClaim = errors.New("invalid audience claim")
+)
+
 type AuthClaims mjwt.BaseTypeClaims[auth.AccessTokenClaims]
 
 type AuthCallback func(rw http.ResponseWriter, req *http.Request, params httprouter.Params, b AuthClaims)
 
-type authChecker struct {
+// AuthChecker validates the bearer token against a mjwt.Verifier and returns an
-	verify mjwt.Verifier
+// error message or continues to the next handler
-	aud    string
+type AuthChecker struct {
-	cb     AuthCallback
+	Verify mjwt.Verifier
+	Aud    string
 }
 
-func (a *authChecker) Handle(rw http.ResponseWriter, req *http.Request, params httprouter.Params) {
+// Middleware is a httprouter.Handle layer to authenticate requests
-	// Get bearer token
+func (a *AuthChecker) Middleware(cb AuthCallback) httprouter.Handle {
-	bearer := utils.GetBearer(req)
+	return func(rw http.ResponseWriter, req *http.Request, params httprouter.Params) {
-	if bearer == "" {
+		// Get bearer token
-		apiError(rw, http.StatusForbidden, "Missing bearer token")
+		bearer := utils.GetBearer(req)
-		return
+		if bearer == "" {
-	}
+			apiError(rw, http.StatusForbidden, "Missing bearer token")
+			return
+		}
 
+		b, err := a.Check(bearer)
+		switch {
+		case errors.Is(err, ErrInvalidToken):
+			apiError(rw, http.StatusForbidden, "Invalid token")
+			return
+		case errors.Is(err, ErrInvalidAudClaim):
+			apiError(rw, http.StatusForbidden, "Invalid audience claim")
+			return
+		case err != nil:
+			apiError(rw, http.StatusForbidden, "Unknown error")
+			return
+		}
+
+		cb(rw, req, params, b)
+	}
+}
+
+// Check takes a token and validates whether it is verified and contains the
+// correct audience claim
+func (a *AuthChecker) Check(token string) (AuthClaims, error) {
 	// Read claims from mjwt
-	_, b, err := mjwt.ExtractClaims[auth.AccessTokenClaims](a.verify, bearer)
+	_, b, err := mjwt.ExtractClaims[auth.AccessTokenClaims](a.Verify, token)
 	if err != nil {
-		apiError(rw, http.StatusForbidden, "Invalid token")
+		return AuthClaims{}, ErrInvalidToken
-		return
 	}
 
+	// Check aud value
 	var validAud bool
 	for _, i := range b.Audience {
-		if subtle.ConstantTimeCompare([]byte(i), []byte(a.aud)) == 1 {
+		if subtle.ConstantTimeCompare([]byte(i), []byte(a.Aud)) == 1 {
 			validAud = true
 		}
 	}
 	if !validAud {
-		apiError(rw, http.StatusForbidden, "Invalid audience claim")
+		return AuthClaims{}, ErrInvalidAudClaim
-		return
 	}
 
-	a.cb(rw, req, params, AuthClaims(b))
+	return AuthClaims(b), nil
-}
-
-// CheckAuth validates the bearer token against a mjwt.Verifier and returns an
-// error message or continues to the next handler
-func CheckAuth(verify mjwt.Verifier, aud string) func(cb AuthCallback) httprouter.Handle {
-	return func(cb AuthCallback) httprouter.Handle {
-		return (&authChecker{
-			verify: verify,
-			aud:    aud,
-			cb:     cb,
-		}).Handle
-	}
 }

cmd/lotus/main.go

@@ -47,7 +47,7 @@ func main() {
 		log.Fatalf("[Lotus] Failed to load MJWT verifier public key from file '%s': %s", filepath.Join(wd, "signer.public.pem"), err)
 	}
 
-	userAuth := api.CheckAuth(verify, conf.Audience)
+	userAuth := &api.AuthChecker{Verify: verify, Aud: conf.Audience}
 	srv := api.SetupApiServer(conf.Listen, userAuth, &conf.SendMail, &conf.Imap)
 	log.Printf("[Lotus] Starting API server on: '%s'\n", srv.Addr)
 	go utils.RunBackgroundHttp("Lotus", srv)

go.mod

@@ -10,7 +10,8 @@ require (
 	github.com/emersion/go-message v0.16.0
 	github.com/emersion/go-sasl v0.0.0-20220912192320-0145f2c60ead
 	github.com/emersion/go-smtp v0.17.0
-	github.com/go-sql-driver/mysql v1.7.1
+	github.com/golang-jwt/jwt/v4 v4.5.0
+	github.com/gorilla/websocket v1.5.0
 	github.com/hydrogen18/memlistener v1.0.0
 	github.com/julienschmidt/httprouter v1.3.0
 	github.com/stretchr/testify v1.8.4
@@ -20,7 +21,6 @@ require (
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/emersion/go-textwrapper v0.0.0-20200911093747-65d896831594 // indirect
-	github.com/golang-jwt/jwt/v4 v4.5.0 // indirect
 	github.com/kr/text v0.2.0 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect

go.sum

@@ -19,10 +19,10 @@ github.com/emersion/go-smtp v0.17.0 h1:tq90evlrcyqRfE6DSXaWVH54oX6OuZOQECEmhWBME
 github.com/emersion/go-smtp v0.17.0/go.mod h1:qm27SGYgoIPRot6ubfQ/GpiPy/g3PaZAVRxiO/sDUgQ=
 github.com/emersion/go-textwrapper v0.0.0-20200911093747-65d896831594 h1:IbFBtwoTQyw0fIM5xv1HF+Y+3ZijDR839WMulgxCcUY=
 github.com/emersion/go-textwrapper v0.0.0-20200911093747-65d896831594/go.mod h1:aqO8z8wPrjkscevZJFVE1wXJrLpC5LtJG7fqLOsPb2U=
-github.com/go-sql-driver/mysql v1.7.1 h1:lUIinVbN1DY0xBg0eMOzmmtGoHwWBbvnWubQUrtU8EI=
-github.com/go-sql-driver/mysql v1.7.1/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI=
 github.com/golang-jwt/jwt/v4 v4.5.0 h1:7cYmW1XlMY7h7ii7UhUyChSgS5wUJEnm9uZVTGqOWzg=
 github.com/golang-jwt/jwt/v4 v4.5.0/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0=
+github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=
+github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/hydrogen18/memlistener v1.0.0 h1:JR7eDj8HD6eXrc5fWLbSUnfcQFL06PYvCc0DKQnWfaU=
 github.com/hydrogen18/memlistener v1.0.0/go.mod h1:qEIFzExnS6016fRpRfxrExeVn2gbClQA99gQhnIcdhE=
 github.com/julienschmidt/httprouter v1.3.0 h1:U0609e9tgbseu3rBINet9P48AI/D3oJs4dN7jwJOQ1U=

imap/client.go

@@ -18,6 +18,44 @@ type Client struct {
 	ic *client.Client
 }
 
+func (c *Client) HandleWS(action string, args []string) (map[string]any, error) {
+	switch action {
+	case "copy":
+		// TODO: implementation
+	case "create":
+		// TODO: implementation
+	case "delete":
+		// TODO: implementation
+	case "select":
+		// TODO: implementation
+	case "fetch":
+		// TODO: implementation
+	case "list":
+		a := make([]*imap.MailboxInfo, 0)
+		b := make(chan *imap.MailboxInfo, 10)
+		go func() {
+			for info := range b {
+				a = append(a, info)
+			}
+		}()
+		err := c.ic.List(args[0], args[1], b)
+		if err != nil {
+			return nil, err
+		}
+		return map[string]any{"Info": a}, nil
+	case "move":
+		// TODO: implementation
+	case "rename":
+		// TODO: implementation
+	case "search":
+		// TODO: implementation
+	case "status":
+		// TODO: implementation
+	}
+	_ = args
+	return map[string]any{"Error": "Not implemented"}, nil
+}
+
 func (c *Client) Append(name string, flags []string, date time.Time, msg imap.Literal) error {
 	return c.ic.Append(name, flags, date, msg)
 }