Skip to content

Instantly share code, notes, and snippets.

@boskiv

boskiv/websocket.go

Created April 22, 2023 13:16
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save boskiv/aa5e8da0941f83f588ac29da2a58b22d to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save boskiv/aa5e8da0941f83f588ac29da2a58b22d to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. boskiv created this gist Apr 22, 2023.
    88 changes: 88 additions & 0 deletions websocket.go
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,88 @@
    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
    }
    }
    }