Go - websocket

 

 

server.go:

package main

import (
    "fmt"
    "log"
    "net/http"
    "os"
    "time"

    "github.com/gorilla/websocket"
)

var port = ":1234"
var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}

func rootHandler(w http.ResponseWriter, r *http.Request) {
    fmt.Fprintln(w, "Welcome!")
    fmt.Fprintln(w, "Please use /ws for WebSocket!")
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
    log.Println("Connection from:", r.Host)
    ws, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        fmt.Println("upgrader.Upgrade:", err)
        return
    }
    defer ws.Close()

    for {
        mt, message, err := ws.ReadMessage()
        if err != nil {
            fmt.Println("From", r.Host, "read", err)
            break
        }
        log.Print("Received:", string(message))

        err = ws.WriteMessage(mt, message)
        if err != nil {
            fmt.Println("WriteMessage:", err)
            break
        }
    }
}

func main() {
    if len(os.Args) != 1 {
        port = ":" + os.Args[1]
    }

    mux := http.NewServeMux()
    s := &http.Server{
        Addr:         port,
        Handler:      mux,
        IdleTimeout:  10 * time.Second,
        ReadTimeout:  time.Second,
        WriteTimeout: time.Second,
    }
    mux.Handle("/", http.HandlerFunc(rootHandler))
    mux.Handle("/ws", http.HandlerFunc(wsHandler))
    log.Println("Listening to TCP port", port)
    err := s.ListenAndServe()
    if err != nil {
        log.Panic(err)
    }
}

 

zzh@ZZHPC:/zdata/Github/ztest/server$ go run server.go
2024/06/13 22:20:16 Listening to TCP port :1234
2024/06/13 22:20:29 Connection from: localhost:1234
2024/06/13 22:20:41 Received:Hello from websocat!
2024/06/13 22:20:47 Received:Bye!
From localhost:1234 read websocket: close 1005 (no status)

 

zzh@ZZHPC:~$ websocat ws://localhost:1234/ws
Hello from websocat!
Hello from websocat!
Bye!
Bye!
ctrl+D

 

client.go:

package main

import (
    "bufio"
    "fmt"
    "log"
    "net/url"
    "os"
    "os/signal"
    "syscall"
    "time"

    "github.com/gorilla/websocket"
)

var (
    SERVER = ""
    PATH = ""
    TIMESWAIT = 0
    TIMESWAITMAX = 5
    in = bufio.NewReader(os.Stdin)
)

func getInput(input chan string) {
    result, err := in.ReadString('\n')
    if err != nil {
        log.Println(err)
        return
    }
    input <- result
}

func main() {
    args := os.Args
    if len(args) != 3 {
        fmt.Println("Need SERVER + PATH!")
        return
    }

    SERVER = args[1]
    PATH = args[2]
    fmt.Println("Connecting to:", SERVER, "at", PATH)

    interrupt := make(chan os.Signal, 1)
    signal.Notify(interrupt, os.Interrupt)

    input := make(chan string, 1)
    go getInput(input)

    URL := url.URL{Scheme: "ws", Host: SERVER, Path: PATH}
    conn, _, err := websocket.DefaultDialer.Dial(URL.String(), nil)
    if err != nil {
        log.Println("Error:", err)
        return
    }
    defer conn.Close()

    done := make(chan struct{})
    go func() {
        defer close(done)
        for {
            _, message, err := conn.ReadMessage()
            if err != nil {
                log.Println("ReadMessage() error:", err)
                return
            }
            log.Printf("Received: %s", message)
        }
    }()

    for {
        select {
        case <-time.After(10 * time.Second):
            log.Println("Please give me input!", TIMESWAIT)
            TIMESWAIT++
            if TIMESWAIT > TIMESWAITMAX {
                syscall.Kill(syscall.Getpid(), syscall.SIGINT)
            }
        case <-done:
            return
        case t := <-input:
            err := conn.WriteMessage(websocket.TextMessage, []byte(t))
            if err != nil {
                log.Println("Write error:", err)
                return
            }
            TIMESWAIT = 0
            go getInput(input)
        case <-interrupt:
            log.Println("Caught interrupt signal - quitting!")
            err := conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
            if err != nil {
                log.Println("Write close error:", err)
                return
            }
            select {
            case <-done:
            case <-time.After(2 * time.Second):
            }
            return
        }
    }
}

 

zzh@ZZHPC:/zdata/Github/ztest/server$ go run server.go
2024/06/14 15:30:02 Listening to TCP port :1234
2024/06/14 15:30:10 Connection from: localhost:1234
2024/06/14 15:30:14 Received:Hello there!
From localhost:1234 read websocket: close 1000 (normal)

 

zzh@ZZHPC:/zdata/Github/ztest/client$ go run client.go localhost:1234 ws
Connecting to: localhost:1234 at ws
Hello there!
2024/06/14 15:30:14 Received: Hello there!
2024/06/14 15:30:24 Please give me input! 0
2024/06/14 15:30:34 Please give me input! 1
^C2024/06/14 15:30:37 Caught interrupt signal - quitting!
2024/06/14 15:30:37 ReadMessage() error: websocket: close 1000 (normal)