websocket async encription

websocket.go

package main

import (
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"log"
	"net/http"
	"sync"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

type UserKey struct {
	sync.RWMutex
	Key []byte
}

func main() {
	userKeys := make(map[string]*UserKey)

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// Upgrade the HTTP connection to a WebSocket connection
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Println(err)
			return
		}
		defer conn.Close()

		// Get the user key from the HTTP headers
		userKeyStr := r.Header.Get("X-User-Key")
		if userKeyStr == "" {
			log.Println("User key not found")
			return
		}
		userKey := []byte(userKeyStr)

		// Create or get the user key from the userKeys map
		var mu sync.Mutex
		key, ok := userKeys[userKeyStr]
		if !ok {
			mu.Lock()
			key = &UserKey{}
			userKeys[userKeyStr] = key
			mu.Unlock()
		}

		// Generate a new AES cipher block using the user key
		block, err := aes.NewCipher(userKey)
		if err != nil {
			log.Println(err)
			return
		}

		// Generate a new GCM AEAD cipher using the block
		aead, err := cipher.NewGCM(block)
		if err != nil {
			log.Println(err)
			return
		}

		// Read and write messages on the WebSocket connection
		for {
			// Read a message from the WebSocket connection
			_, msg, err := conn.ReadMessage()
			if err != nil {
				log.Println(err)
				break
			}

			// Decrypt the message using the AEAD cipher
			nonceSize := aead.NonceSize()
			if len(msg) < nonceSize {
				log.Println("Invalid message length")
				break
			}
			nonce, ciphertext := msg[:nonceSize], msg[nonceSize:]
			plaintext, err := aead.Open(nil, nonce, ciphertext, nil)
			if err != nil {
				log.Println(err)
				break
			}
			log.Printf("Received message: %s", plaintext)

			// Encrypt a message using the AEAD cipher
			nonce = make([]byte, nonceSize)
			if _, err := rand.Read(nonce); err != nil {
				log.Println(err)
				break
			}
			ciphertext = aead.Seal(nonce, nonce, plaintext, nil)
			msg = append(nonce, ciphertext...)
			log.Printf("Sending message: %s", msg)

			// Write the encrypted message to the WebSocket connection
			err = conn.WriteMessage(websocket.TextMessage, msg)
			if err != nil {
				log.Println(err)
				break
			}
		}

		// Delete the user key from the userKeys map when the connection is closed
		key.Lock()
		delete(userKeys, userKeyStr)
		key.Unlock()
	})

	log.Fatal(http.ListenAndServe(":8080", nil))
}