Skip to content

Instantly share code, notes, and snippets.

@boskiv
Created April 22, 2023 13:16
Show Gist options
  • Save boskiv/aa5e8da0941f83f588ac29da2a58b22d to your computer and use it in GitHub Desktop.
Save boskiv/aa5e8da0941f83f588ac29da2a58b22d to your computer and use it in GitHub Desktop.
Websocket with encryption
package main
import (
"crypto/aes"
"crypto/cipher"
"encoding/base64"
"log"
"net/http"
"github.com/gorilla/websocket"
)
var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
}
var aesKey = []byte("0123456789012345") // 16 bytes AES key
func main() {
http.HandleFunc("/", handleWebSocket)
log.Fatal(http.ListenAndServe(":8080", nil))
}
func handleWebSocket(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()
// Generate a new AES cipher block using the key
block, err := aes.NewCipher(aesKey)
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
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment