From 01662618f725822fc1c1dd16ac2a7f41dc2070ce Mon Sep 17 00:00:00 2001 From: Tom Wiesing Date: Mon, 22 Apr 2024 01:57:48 +0200 Subject: [PATCH] Use pkglib instead of custom packages --- go.mod | 3 +- go.sum | 4 + pkg/password/password.go | 44 --------- pkg/repo/uploadable.go | 70 ++++++++------- pkg/wshandler/handler.go | 188 --------------------------------------- pkg/wshandler/pool.go | 22 ----- 6 files changed, 46 insertions(+), 285 deletions(-) delete mode 100644 pkg/password/password.go delete mode 100644 pkg/wshandler/handler.go delete mode 100644 pkg/wshandler/pool.go diff --git a/go.mod b/go.mod index f3929ae..03870ca 100644 --- a/go.mod +++ b/go.mod @@ -5,16 +5,17 @@ go 1.22.2 require ( github.com/die-net/lrucache v0.0.0-20220628165024-20a71bc65bf1 github.com/google/go-github v17.0.0+incompatible - github.com/gorilla/websocket v1.5.1 github.com/gregjones/httpcache v0.0.0-20190611155906-901d90724c79 github.com/mpolden/echoip v0.0.0-20230521182614-d84665c26cf7 github.com/pkg/errors v0.9.1 + github.com/tkw1536/pkglib v0.0.0-20240421230152-7fcf00edc7df golang.org/x/crypto v0.22.0 golang.org/x/oauth2 v0.19.0 ) require ( github.com/google/go-querystring v1.1.0 // indirect + github.com/gorilla/websocket v1.5.1 // indirect golang.org/x/net v0.24.0 // indirect golang.org/x/sys v0.19.0 // indirect ) diff --git a/go.sum b/go.sum index cc57806..83c3eaa 100644 --- a/go.sum +++ b/go.sum @@ -29,6 +29,10 @@ github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.5 h1:s5PTfem8p8EbKQOctVV53k6jCJt3UX4IEJzwh+C324Q= github.com/stretchr/testify v1.7.5/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/tkw1536/pkglib v0.0.0-20240421230152-7fcf00edc7df h1:ggQcsE68zLKjYf+kdQG10wK9+KyV2PO9wpBKzkX0wEw= +github.com/tkw1536/pkglib v0.0.0-20240421230152-7fcf00edc7df/go.mod h1:P/9GGxNGvEZsYShL7bv57UdiHkR0l6AWz1IDMVJGzlY= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= golang.org/x/crypto v0.22.0 h1:g1v0xeRhjcugydODzvb3mEM9SQ0HGp9s/nh3COQ/C30= golang.org/x/crypto v0.22.0/go.mod h1:vr6Su+7cTlO45qkww3VDJlzDn0ctJvRgYbC2NvXHt+M= golang.org/x/net v0.24.0 h1:1PcaxkF854Fu3+lvBIx5SYn9wRlBzzcnHZSiaFFAb0w= diff --git a/pkg/password/password.go b/pkg/password/password.go deleted file mode 100644 index 4db28ef..0000000 --- a/pkg/password/password.go +++ /dev/null @@ -1,44 +0,0 @@ -package password - -import ( - "crypto/rand" - "math/big" - "strings" -) - -const PasswordCharSet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" - -var passwordCharCount = big.NewInt(int64(len(PasswordCharSet))) - -// Password returns a randomly generated string with the provided length. -// It consists of alphanumeric characters only. -// -// When an error occurs, it is guaranteed to return "", err. -// [rand.Reader] is used as the source of randomness. -func Password(length int) (string, error) { - if length < 0 { - panic("length < 0") - } - - // create a buffer to write the string to! - var password strings.Builder - password.Grow(length) - - for i := 0; i < length; i++ { - - // grab a random bIndex! - bIndex, err := rand.Int(rand.Reader, passwordCharCount) - if err != nil { - return "", err - } - - // and use that index! - index := int(bIndex.Int64()) - if err := password.WriteByte(PasswordCharSet[index]); err != nil { - return "", err - } - } - - // return the password! - return password.String(), nil -} diff --git a/pkg/repo/uploadable.go b/pkg/repo/uploadable.go index af9763a..3d3d0b9 100644 --- a/pkg/repo/uploadable.go +++ b/pkg/repo/uploadable.go @@ -2,6 +2,7 @@ package repo import ( "context" + "crypto/rand" "fmt" "io" "net/http" @@ -9,12 +10,14 @@ import ( "sync/atomic" "github.com/pkg/errors" - "github.com/tkw1536/akhttpd/pkg/password" - "github.com/tkw1536/akhttpd/pkg/wshandler" "golang.org/x/crypto/ssh" _ "embed" + + "github.com/tkw1536/pkglib/lazy" + "github.com/tkw1536/pkglib/password" + "github.com/tkw1536/pkglib/websocketx" ) // spellchecker:words akhttpd wshandler userkeys @@ -30,6 +33,8 @@ type UploadableKeys struct { lock sync.RWMutex data map[string][]ssh.PublicKey + + server lazy.Lazy[*websocketx.Server] } var errUserKeysNotConfigured = UserNotFoundError{errors.New("User is not configured in UserKeys")} @@ -80,13 +85,26 @@ func (uk *UploadableKeys) Register(keys ...ssh.PublicKey) (username string, clea // username generates a new username func (uk *UploadableKeys) username() string { - hash, err := password.Password(10) + hash, err := password.Generate(rand.Reader, 10, password.DefaultCharSet) if err != nil { // fallback to a counter-based approach return fmt.Sprintf("%s%d", uk.Prefix, atomic.AddUint64(&uk.counter, 1)) } return uk.Prefix + hash } +func (uk *UploadableKeys) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if !uk.auth(w, r) { + return + } + + uk.server.Get(func() *websocketx.Server { + return &websocketx.Server{ + Handler: uk.handleWS, + Fallback: http.HandlerFunc(uk.handleHTTP), + } + }).ServeHTTP(w, r) +} + //go:embed uploadable.min.html var uploadableHTML []byte @@ -112,39 +130,31 @@ func (uk *UploadableKeys) auth(w http.ResponseWriter, r *http.Request) bool { return false } -func (uk *UploadableKeys) ServeHTTP(w http.ResponseWriter, r *http.Request) { - if !uk.auth(w, r) { +func (uk *UploadableKeys) handleWS(conn *websocketx.Connection) { + key, ok := <-conn.Read() + if !ok { return } - // if an upgrade to the websocket was requested, serve a websocket! - if r.Header.Get("Upgrade") == "websocket" { - wshandler.Handle(w, r, func(messenger wshandler.WebSocket) { - key, ok := messenger.Read() // wait for any kind of message - if !ok { - return - } - - // read a private key from the connection! - pk, _, _, _, err := ssh.ParseAuthorizedKey([]byte(key)) - if err != nil { - return - } - - // register the key - username, cleanup := uk.Register(pk) - defer cleanup() - - // write the username back! - if !messenger.Write(username) { - return - } - - messenger.Wait() - }) + // read a private key from the connection! + pk, _, _, _, err := ssh.ParseAuthorizedKey(key.Body) + if err != nil { return } + // register the key + username, cleanup := uk.Register(pk) + defer cleanup() + + // Write the username back + <-conn.WriteText(username) + + // and wait for the connection to be closed + // by the client + <-conn.Context().Done() +} + +func (uk *UploadableKeys) handleHTTP(w http.ResponseWriter, r *http.Request) { w.Header().Add("Content-Type", "text/html") w.WriteHeader(http.StatusOK) w.Write(uploadableHTML) diff --git a/pkg/wshandler/handler.go b/pkg/wshandler/handler.go deleted file mode 100644 index eecf64a..0000000 --- a/pkg/wshandler/handler.go +++ /dev/null @@ -1,188 +0,0 @@ -package wshandler - -import ( - "context" - "net/http" - "sync" - "time" - - "github.com/gorilla/websocket" -) - -// WebSocket represents a simplified connection -type WebSocket interface { - // Wait waits for the connection to close - Wait() - - // Read reads a message - Read() (string, bool) - - // Write writes a message - Write(string) bool -} - -type HandleFunc func(WebSocket) - -var upgrader websocket.Upgrader - -// Handler handles WebSocket connections with a simple messages -type Handler struct { - conn *websocket.Conn // underlying connection - context context.Context // context to cancel the connection - cancel context.CancelFunc - - wg sync.WaitGroup // blocks all the ongoing tasks - - // incoming and outgoing tasks - incoming chan string - outgoing chan string - - // HandleFunc is called to implement the application logic. - // It should return to terminate a connection. - HandleFunc HandleFunc -} - -const ( - writeWait = 10 * time.Second - pongWait = time.Minute - pingInterval = (pongWait * 9) / 10 - maxMessageSize = 2048 -) - -func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { - defer h.wg.Wait() // wait for everything to finish before returning to the caller - - var err error - h.conn, err = upgrader.Upgrade(w, r, nil) - if err != nil { - return - } - - // wait for the context to be cancelled, then close the connection - h.wg.Add(1) - go func() { - defer h.wg.Done() - <-h.context.Done() - h.conn.Close() - }() - - // start receiving and sending messages - h.wg.Add(2) - go h.sendMessages() - go h.recvMessages() - - // start the application logic - h.wg.Add(1) - go h.handle() -} - -func (h *Handler) handle() { - defer func() { - h.wg.Done() - h.cancel() - }() - - h.HandleFunc(h) -} - -func (h *Handler) sendMessages() { - // close connection when done! - defer func() { - h.wg.Done() - h.cancel() - }() - - // setup a timer for pings! - ticker := time.NewTicker(pingInterval) - defer ticker.Stop() - - for { - select { - // everything is done! - case <-h.context.Done(): - return - - // send outgoing messages - case content := <-h.outgoing: - if err := h.writeRaw(websocket.TextMessage, []byte(content)); err != nil { - return - } - // send a ping message - case <-ticker.C: - if err := h.writeRaw(websocket.PingMessage, []byte{}); err != nil { - return - } - } - } -} - -// writeRaw writes to the underlying socket -func (h *Handler) writeRaw(messageType int, data []byte) error { - h.conn.SetWriteDeadline(time.Now().Add(writeWait)) - return h.conn.WriteMessage(messageType, data) -} - -// Read reads a message sent from the client -func (h *Handler) Read() (message string, ok bool) { - select { - case message = <-h.incoming: - return message, true - case <-h.context.Done(): - return "", false - } -} - -func (h *Handler) Wait() { - <-h.context.Done() -} - -// Write writes a message to the client -func (sh *Handler) Write(message string) bool { - select { - case sh.outgoing <- message: - return true - case <-sh.context.Done(): - return false - } -} - -func (h *Handler) recvMessages() { - // close connection when done! - defer func() { - h.wg.Done() - h.cancel() - }() - - h.conn.SetReadLimit(maxMessageSize) - - // configure a pong handler - h.conn.SetReadDeadline(time.Now().Add(pongWait)) - h.conn.SetPongHandler(func(string) error { h.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil }) - - // handle incoming messages - for { - _, message, err := h.conn.ReadMessage() - if err != nil { - return - } - - // send a message to the incoming channel, or cancel - select { - case h.incoming <- string(message): - case <-h.context.Done(): - return - } - - } -} - -// Reset resets this SocketHandler to default -func (h *Handler) Reset(HandleFunc HandleFunc) { - h.HandleFunc = HandleFunc - - h.conn = nil - - h.incoming = make(chan string) - h.outgoing = make(chan string) - h.context, h.cancel = context.WithCancel(context.Background()) -} diff --git a/pkg/wshandler/pool.go b/pkg/wshandler/pool.go deleted file mode 100644 index 69554b1..0000000 --- a/pkg/wshandler/pool.go +++ /dev/null @@ -1,22 +0,0 @@ -package wshandler - -import ( - "net/http" - "sync" -) - -var pool = sync.Pool{ - New: func() any { - return new(Handler) - }, -} - -// Handle instantiates a new handler, and calls it with the provided -func Handle(w http.ResponseWriter, r *http.Request, HandleFunc HandleFunc) { - handler := pool.Get().(*Handler) - defer pool.Put(handler) - - handler.Reset(HandleFunc) - handler.ServeHTTP(w, r) - handler.Reset(nil) -}