Use pkglib instead of custom packages

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
 )

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=

pkg/repo/uploadable.go

@@ -2,19 +2,22 @@ package repo
 
 import (
 	"context"
+	"crypto/rand"
 	"fmt"
 	"io"
 	"net/http"
 	"sync"
 	"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)