feat: connect terminal service to backend

feat: put terminal service into container

refactor: separate main into terminal and websocket

feat: create zsh using commander in home

feat: pass backendUrl to container

feat: pass token to ws server to cmd-hook

feat: gen terminal access token for frontend

feat: success to connect through traefik

fix: handle invalid jwt

feat: connect to session container via traefik

app/src/components/game/QuestPart.vue

@@ -22,6 +22,7 @@ const tasks = computed(() => {
   <div class="h-full w-full max-h-full">
     <p id="topic" class="text-text font-black" v-if="sessionStore.quest">
       {{ sessionStore.quest.title }}
+      {{ sessionStore.session?.token }}
     </p>
     <br />
     <div>

app/src/store/session.ts

@@ -94,12 +94,12 @@ const useSession = defineStore('session', {
     },
     setup() {
       if (hasSetup) return
-      socket.on('terminal', (data: string) => {
-        term.write(data)
-      })
-      term.onData((data: string) => {
-        socket.emit('terminal', data)
-      })
+      // socket.on('terminal', (data: string) => {
+      //   term.write(data)
+      // })
+      // term.onData((data: string) => {
+      //   socket.emit('terminal', data)
+      // })
       socket.on('graph', (event: FileGraphUpdateEvent) => {
         if (!this.session) return
         this.session.graph.handleEvent(event)

app/src/types.ts

@@ -21,4 +21,5 @@ export interface Session {
   pwd: string
   tasks: ITask[]
   responses: IResponse[]
+  token: string
 }

docker-compose.yml

@@ -1,5 +1,22 @@
 version: '3'
 services:
+  traefik:
+    image: traefik:v3.3
+    networks:
+      - players-dev
+    command:
+      - '--api.insecure=true'
+      - '--providers.docker=true'
+      - '--providers.docker.exposedbydefault=false'
+      - '--entrypoints.web.address=:80'
+      - '--entrypoints.websecure.address=:443'
+    ports:
+      - '80:80'
+      - '443:443'
+      - '8080:8080'
+    volumes:
+      - /var/run/docker.sock:/var/run/docker.sock:ro
+
   db:
     image: mongo:6
     restart: always
@@ -19,6 +36,13 @@ services:
       dockerfile: ./quests/Dockerfile
     extra_hosts:
       - 'host.docker.internal:host-gateway'
+    labels:
+      - 'traefik.enable=true'
+      - 'traefik.http.routers.base.rule=PathPrefix(`/terminal/base`)'
+      - 'traefik.http.middlewares.base-stripprefix.stripprefix.prefixes=/terminal/base'
+      - 'traefik.http.routers.base.middlewares=base-stripprefix'
+      - 'traefik.http.routers.base.entrypoints=web'
+      - 'traefik.http.services.base.loadbalancer.server.port=9090'
 
 networks:
   internal:

quests/.includes/zshrc

@@ -29,7 +29,7 @@ precmd() {
     export CMD_ERROR_FILE=/tmp/cmd.error
 
     # Now run your Node.js script with the result as an argument
-    /usr/local/bin/cmd-hook; rm /tmp/cmd.out /tmp/cmd.error 2>/dev/null
+    /usr/local/lib/cmd-hook/cmd-hook; rm /tmp/cmd.out /tmp/cmd.error 2>/dev/null
 
     # Reset variables
     unset CMD_START_TIME

quests/Dockerfile

@@ -1,15 +1,18 @@
-FROM golang:1.24-alpine as builder
-
+FROM golang:1.24-alpine as cmd-hook-builder
 WORKDIR /app
-
 COPY ./services/cmd-hook/go.mod ./
-
 RUN go mod download
-
 COPY ./services/cmd-hook ./
-
 RUN go build -o cmd-hook
 
+FROM golang:1.24-alpine as terminal-builder
+WORKDIR /app
+COPY ./services/terminal-service/go.mod ./
+COPY ./services/terminal-service/go.sum ./
+RUN go mod download
+COPY ./services/terminal-service ./
+RUN go build -o terminal-service
+
 FROM debian:bookworm-slim as main
 
 LABEL maintainer="wancat@wancat.cc"
@@ -18,16 +21,13 @@ RUN apt-get update && apt-get install -y \
   vim \
   zsh \
   curl \
-  openssh-server
-
-COPY --from=builder /app/cmd-hook /usr/local/bin/cmd-hook
+  sudo
 COPY ./quests/.includes/zshrc /etc/zsh/zshrc
-RUN rm -f /etc/motd /etc/issue /etc/issue.net /etc/update-motd.d/*
 RUN rm -rf /etc/skel && mkdir -p /etc/skel
-RUN mkdir -p /etc/ssh/authorized_keys && chmod 755 /etc/ssh/authorized_keys
-COPY ./quests/.includes/sshd_config /etc/ssh/sshd_config
 COPY ./quests/entrypoint.sh /entrypoint.sh
+COPY --from=cmd-hook-builder /app/cmd-hook /usr/local/lib/cmd-hook/cmd-hook
+COPY --from=terminal-builder /app/terminal-service /usr/local/lib/terminal-service/terminal-service
 
-EXPOSE 22
+EXPOSE 9090
 
 ENTRYPOINT ["/bin/sh", "/entrypoint.sh"]

quests/entrypoint.sh

@@ -3,19 +3,16 @@
 USERNAME=commander
 
 echo "Initializing container for $USERNAME"
+echo "BACKEND_URL: $BACKEND_URL"
 
 if ! id -u ${USERNAME} >/dev/null 2>&1; then
   echo "Adding user $USERNAME"
   adduser --disabled-password --gecos '' --shell /bin/zsh ${USERNAME} && \
   find / -name ".gitkeep" -type f -exec rm -f {} +
-
-  echo "Adding public key for $USERNAME"
-  echo "Public key: $PUB_KEY"
-  PUB_KEY_PATH=/etc/ssh/authorized_keys/${USERNAME}
-  echo $PUB_KEY > ${PUB_KEY_PATH}
-  chmod 644 ${PUB_KEY_PATH}
 fi
 
-mkdir -p /run/sshd
-echo "Starting SSH server"
-/usr/sbin/sshd -D -e
+# Preserve environment variables
+export BACKEND_URL
+
+echo "Starting terminal service"
+su - ${USERNAME} -w BACKEND_URL -c "/usr/local/lib/terminal-service/terminal-service"

server/src/api/controllers/commandController.ts

@@ -7,11 +7,29 @@ import { asyncHandler } from '../../middleware/error'
 import { questManager } from '../../models/quest'
 import { CLIFileExistenceChecker } from '../../containers/cli'
 import config from '../../config'
+import { verifyJWT } from '../../utils/auth'
 
 export const newCommand = asyncHandler(async (req: Request, res: Response) => {
   if (!config.isProduction) {
     console.log('newCommand', JSON.stringify(req.body, null, 2))
   }
+  const token = req.headers.authorization?.split(' ')[1]
+  if (!token) {
+    res.status(401).json({ message: 'Unauthorized' })
+    return
+  }
+  let sessionId: string
+  try {
+    const decoded = await verifyJWT(token)
+    sessionId = decoded.sessionId
+    if (!sessionId) {
+      res.status(401).json({ message: 'Unauthorized' })
+      return
+    }
+  } catch (error) {
+    res.status(401).json({ message: 'Unauthorized' })
+    return
+  }
   const body = commandSchema.safeParse(req.body)
   if (!body.success) {
     res.status(400).json({ message: 'Invalid command', errors: body.error })
@@ -19,7 +37,6 @@ export const newCommand = asyncHandler(async (req: Request, res: Response) => {
   }
   const command = body.data
 
-  const { sessionId } = req.user as any
   const session = await Session.findById(sessionId)
   if (!session) {
     res.status(404).json({ message: 'session not found' })

server/src/api/routes/index.ts

@@ -17,12 +17,7 @@ if (!config.isProduction) {
   })
 }
 
-router.use(
-  '/commands',
-  passport.authenticate('jwt', { session: false }),
-  authRequired,
-  commands
-)
+router.use('/commands', commands)
 router.use('/auth', authRouter)
 
 router.get('/leaderboard', leaderboardController)

server/src/api/socket.ts

@@ -1,12 +1,9 @@
 import { Server } from 'socket.io'
 import type { Socket } from 'socket.io'
-import { Duplex } from 'stream'
 import validator from 'validator'
 import type { Server as HttpServer } from 'http'
 import { ISession, Session } from '../../../packages/models'
-import { getAndStartContainer, attachContainer } from '../containers/docker.js'
 import SessionMiddleware from '../middleware/session.js'
-import { genJWT } from '../utils/auth.js'
 import logger from '../utils/logger.js'
 
 type SocketCallback = (event: string, ...args: any[]) => void
@@ -61,41 +58,16 @@ async function connectContainer(socket: Socket, next: (err?: Error) => void) {
     return
   }
 
-  const token = await genJWT({
-    sessionId: session.id,
-  })
-
-  if (!session.containerId) {
-    next(new Error('Session has no container.'))
-    return
-  }
-
-  try {
-    const container = await getAndStartContainer(session.containerId)
-    const stream = await attachContainer(container, { token })
-    socket.data.context = {
-      session,
-      stream,
-    }
-    next()
-  } catch (err) {
-    logger.error(err)
-    next(new Error('Failed to connect to container.'))
+  socket.data.context = {
+    session,
   }
+  next()
 }
 
 function onConnect(socket: Socket) {
-  const { session, stream } = socket.data.context as {
+  const { session } = socket.data.context as {
     session: ISession
-    stream: Duplex
   }
-  stream.on('data', (chunk: Buffer) => {
-    socket.emit('terminal', chunk.toString())
-  })
-
-  socket.on('terminal', function incoming(message) {
-    stream.write(message)
-  })
 
   const socketCallback = (event: string, ...data: any[]) => {
     socket.emit(event, ...data)
@@ -104,9 +76,7 @@ function onConnect(socket: Socket) {
   listenToSession(session._id.toString(), socketCallback)
 
   socket.on('disconnect', () => {
-    stream.write('exit\n')
     removeFromSession(session._id.toString(), socketCallback)
-    stream.destroy()
   })
 }
 

server/src/containers/docker.ts

@@ -22,9 +22,17 @@ const newContainerOptions = (
     ExtraHosts: ['host.docker.internal:host-gateway'],
   },
   ExposedPorts: {
-    '22/tcp': {},
+    '9090/tcp': {},
+  },
+  Env: [`BACKEND_URL=${config.backendUrl}`],
+  Labels: {
+    'traefik.enable': 'true',
+    [`traefik.http.routers.terminal-${name}.rule`]: `PathPrefix(\`/terminal/${name}\`)`,
+    [`traefik.http.middlewares.terminal-${name}-stripprefix.stripprefix.prefixes`]: `/terminal/${name}`,
+    [`traefik.http.routers.terminal-${name}.middlewares`]: `terminal-${name}-stripprefix`,
+    [`traefik.http.services.terminal-${name}.loadbalancer.server.port`]: '9090',
+    [`traefik.http.routers.terminal-${name}.entrypoints`]: 'web',
   },
-  Env: [`PUB_KEY=${config.docker.keypair.publicKey}`],
 })
 
 export async function createContainer(
@@ -48,7 +56,8 @@ export async function createContainer(
   if (config.docker.mountCLI) {
     logger.info('Mounting CLI', questId)
     binds.push(
-      `${config.docker.hostProjectRoot}/services/cmd-hook/bin:/usr/local/bin:ro`
+      `${config.docker.hostProjectRoot}/services/cmd-hook/bin:/usr/local/lib/cmd-hook:ro`,
+      `${config.docker.hostProjectRoot}/services/terminal-service/bin:/usr/local/lib/terminal-service:ro`
     )
   }
   const option = newContainerOptions(name, imageId, { binds })
@@ -98,6 +107,7 @@ export async function attachContainer(
   const containerIp = (await container.inspect()).NetworkSettings.Networks[
     config.docker.network
   ].IPAddress
+  console.log('containerIp', containerIp)
 
   let error: Error | null = null
   for (let i = 0; i < 10; i++) {

server/src/models/sessionManager.ts

@@ -1,7 +1,11 @@
 import { HydratedDocument } from 'mongoose'
 import { Session, UserProfile } from '../../../packages/models'
 import type { ISession, IUser } from '../../../packages/models'
-import { createContainer, deleteContainer } from '../containers/docker.js'
+import {
+  createContainer,
+  deleteContainer,
+  getAndStartContainer,
+} from '../containers/docker.js'
 import logger from '../utils/logger.js'
 import { questManager } from './quest.js'
 
@@ -48,11 +52,10 @@ export async function createNewSession(
     throw new Error(`UserProfile ${userId} not found`)
   }
 
-  const container = await createContainer(
-    `quest-${quest.id}-${userProfile.user.username}-${Date.now()}`,
-    quest.id,
-    quest.image
-  )
+  const containerName = `quest-${quest.id}-${userProfile.user.username}-${Date.now()}`
+
+  const container = await createContainer(containerName, quest.id, quest.image)
+  await getAndStartContainer(container.id)
 
   const newSession = new Session({
     user: userId,
@@ -60,6 +63,9 @@ export async function createNewSession(
     containerId: container.id,
   })
 
+  console.log('container', container.id)
+  console.log('containerName', `http://localhost/terminal/${containerName}`)
+
   const progress = userProfile.progress.get(quest.id)
   if (!progress) {
     userProfile.progress.set(quest.id, {

server/src/routers/sessionRouter.ts

@@ -6,10 +6,11 @@ import { createNewSession, isQuestUnlocked } from '../models/sessionManager.js'
 import config from '../config.js'
 import { protectedProcedure, router } from '../trpc.js'
 import { questManager } from '../models/quest.js'
+import { genJWT } from '../utils/auth'
 
 export type SessionDetail = Awaited<ReturnType<typeof sessionDetail>>
 
-function sessionDetail(session: ISession) {
+async function sessionDetail(session: ISession) {
   const quest = questManager.get(session.quest)
   if (!quest) {
     throw new TRPCError({
@@ -22,6 +23,10 @@ function sessionDetail(session: ISession) {
     quest,
     new VoidFileExistenceChecker()
   )
+  const token = await genJWT({
+    sessionId: session._id.toString(),
+  })
+
   return {
     _id: session._id.toString(),
     user: session.user.toString(),
@@ -33,6 +38,7 @@ function sessionDetail(session: ISession) {
     responses: gameSession.getResponses(),
     tasks: gameSession.getTasks(),
     graph: gameSession.getGraph(),
+    token,
   }
 }
 

services/cmd-hook/main.go

@@ -48,7 +48,11 @@ func main() {
 
 	payload.Params = params
 
-	apiEndpoint := os.Getenv("API_ENDPOINT")
+	backendUrl := os.Getenv("BACKEND_URL")
+	if backendUrl == "" {
+		fmt.Fprintf(os.Stderr, "BACKEND_URL is not set\n")
+		os.Exit(1)
+	}
 	token := os.Getenv("TOKEN")
 
 	jsonData, err := json.MarshalIndent(payload, "", "  ")
@@ -58,7 +62,7 @@ func main() {
 	}
 
 	fmt.Println(string(jsonData))
-	req, err := http.NewRequest("POST", fmt.Sprintf("%s/api/v1/commands", apiEndpoint), bytes.NewBuffer(jsonData))
+	req, err := http.NewRequest("POST", fmt.Sprintf("%s/api/v1/commands", backendUrl), bytes.NewBuffer(jsonData))
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "Error creating request: %v\n", err)
 		os.Exit(1)

services/terminal-service/.air.toml

@@ -4,8 +4,8 @@ tmp_dir = "tmp"
 
 [build]
   args_bin = []
-  bin = "./tmp/main"
-  cmd = "go build -o ./tmp/main ."
+  bin = "./bin/terminal-service"
+  cmd = "go build -o ./bin/terminal-service ."
   delay = 1000
   exclude_dir = ["assets", "tmp", "vendor", "testdata"]
   exclude_file = []

services/terminal-service/assets.go

@@ -0,0 +1,6 @@
+package main
+
+import _ "embed"
+
+//go:embed index.html
+var indexHTML []byte

services/terminal-service/index.html

@@ -39,9 +39,13 @@
       term.open(document.getElementById('terminal'))
       fitAddon.fit()
 
+      const token = new URLSearchParams(window.location.search).get('token')
+
       // Create WebSocket connection
       const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
-      const ws = new WebSocket(`${protocol}//${window.location.host}/ws`)
+      const ws = new WebSocket(
+        `${protocol}//${window.location.host}${window.location.pathname}/ws?token=${token}`
+      )
 
       // Handle WebSocket connection
       ws.onopen = () => {

services/terminal-service/main.go

@@ -2,136 +2,16 @@ package main
 
 import (
 	"flag"
-	"io"
+	"fmt"
 	"log"
 	"net/http"
 	"os"
-	"os/exec"
-	"time"
-
-	"github.com/creack/pty"
-	"github.com/gorilla/websocket"
 )
 
 var (
-	addr = flag.String("addr", ":8080", "http service address")
+	addr = flag.String("addr", ":9090", "http service address")
 )
 
-const (
-	// Time allowed to write a message to the peer.
-	writeWait = 10 * time.Second
-
-	// Maximum message size allowed from peer.
-	maxMessageSize = 8192
-
-	// Time allowed to read the next pong message from the peer.
-	pongWait = 60 * time.Second
-
-	// Send pings to peer with this period. Must be less than pongWait.
-	pingPeriod = (pongWait * 9) / 10
-
-	// Time to wait before force close on connection.
-	closeGracePeriod = 10 * time.Second
-)
-
-var upgrader = websocket.Upgrader{
-	ReadBufferSize:  1024,
-	WriteBufferSize: 1024,
-	// Allow all origins for now - we'll restrict this in production
-	CheckOrigin: func(r *http.Request) bool {
-		return true
-	},
-}
-
-type terminal struct {
-	cmd *exec.Cmd
-	pty *os.File
-}
-
-func newTerminal() (*terminal, error) {
-	// Start a new shell
-	cmd := exec.Command("bash")
-	ptmx, err := pty.Start(cmd)
-	if err != nil {
-		return nil, err
-	}
-
-	return &terminal{
-		cmd: cmd,
-		pty: ptmx,
-	}, nil
-}
-
-func (t *terminal) Close() error {
-	if t.pty != nil {
-		t.pty.Close()
-	}
-	if t.cmd != nil {
-		t.cmd.Process.Kill()
-		t.cmd.Wait()
-	}
-	return nil
-}
-
-func pumpStdin(ws *websocket.Conn, w io.Writer) {
-	defer ws.Close()
-	ws.SetReadLimit(maxMessageSize)
-	ws.SetReadDeadline(time.Now().Add(pongWait))
-	ws.SetPongHandler(func(string) error { ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
-	for {
-		_, message, err := ws.ReadMessage()
-		if err != nil {
-			break
-		}
-		if _, err := w.Write(message); err != nil {
-			break
-		}
-	}
-}
-
-func pumpStdout(ws *websocket.Conn, r io.Reader, done chan struct{}) {
-	defer func() {
-	}()
-	buf := make([]byte, 1024)
-	for {
-		n, err := r.Read(buf)
-		if err != nil {
-			break
-		}
-		ws.SetWriteDeadline(time.Now().Add(writeWait))
-		if err := ws.WriteMessage(websocket.TextMessage, buf[:n]); err != nil {
-			ws.Close()
-			break
-		}
-	}
-	close(done)
-
-	ws.SetWriteDeadline(time.Now().Add(writeWait))
-	ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
-	time.Sleep(closeGracePeriod)
-	ws.Close()
-}
-
-func ping(ws *websocket.Conn, done chan struct{}) {
-	ticker := time.NewTicker(pingPeriod)
-	defer ticker.Stop()
-	for {
-		select {
-		case <-ticker.C:
-			if err := ws.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(writeWait)); err != nil {
-				log.Println("ping:", err)
-			}
-		case <-done:
-			return
-		}
-	}
-}
-
-func internalError(ws *websocket.Conn, msg string, err error) {
-	log.Println(msg, err)
-	ws.WriteMessage(websocket.TextMessage, []byte("Internal server error."))
-}
-
 func serveWs(w http.ResponseWriter, r *http.Request) {
 	ws, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
@@ -140,8 +20,14 @@ func serveWs(w http.ResponseWriter, r *http.Request) {
 	}
 	defer ws.Close()
 
+	token := r.URL.Query().Get("token")
+	if token == "" {
+		http.Error(w, "Unauthorized", http.StatusUnauthorized)
+		return
+	}
+
 	// Create new terminal
-	term, err := newTerminal()
+	term, err := newTerminal(token)
 	if err != nil {
 		internalError(ws, "start:", err)
 		return
@@ -167,7 +53,15 @@ func serveHome(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
 		return
 	}
-	http.ServeFile(w, r, "index.html")
+	w.Header().Set("Content-Type", "text/html")
+	w.Write(indexHTML)
+}
+
+func serveHealth(w http.ResponseWriter, r *http.Request) {
+	w.WriteHeader(http.StatusOK)
+	for _, e := range os.Environ() {
+		fmt.Fprintln(w, e)
+	}
 }
 
 func main() {
@@ -175,5 +69,7 @@ func main() {
 	log.SetFlags(0)
 	http.HandleFunc("/", serveHome)
 	http.HandleFunc("/ws", serveWs)
+	http.HandleFunc("/health", serveHealth)
+	fmt.Println("Starting terminal service on", *addr)
 	log.Fatal(http.ListenAndServe(*addr, nil))
 }

services/terminal-service/terminal.go

@@ -0,0 +1,40 @@
+package main
+
+import (
+	"os"
+	"os/exec"
+
+	"github.com/creack/pty"
+)
+
+type terminal struct {
+	cmd *exec.Cmd
+	pty *os.File
+}
+
+func newTerminal(token string) (*terminal, error) {
+	// Start a new shell in the home directory
+	cmd := exec.Command("/usr/bin/zsh")
+	cmd.Dir = os.Getenv("HOME")
+	cmd.Env = append(os.Environ(), "ZDOTDIR=/etc/zsh", "TOKEN="+token)
+	ptmx, err := pty.Start(cmd)
+	if err != nil {
+		return nil, err
+	}
+
+	return &terminal{
+		cmd: cmd,
+		pty: ptmx,
+	}, nil
+}
+
+func (t *terminal) Close() error {
+	if t.pty != nil {
+		t.pty.Close()
+	}
+	if t.cmd != nil {
+		t.cmd.Process.Kill()
+		t.cmd.Wait()
+	}
+	return nil
+}

services/terminal-service/websocket.go

@@ -0,0 +1,95 @@
+package main
+
+import (
+	"io"
+	"log"
+	"net/http"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+const (
+	// Time allowed to write a message to the peer.
+	writeWait = 10 * time.Second
+
+	// Maximum message size allowed from peer.
+	maxMessageSize = 8192
+
+	// Time allowed to read the next pong message from the peer.
+	pongWait = 60 * time.Second
+
+	// Send pings to peer with this period. Must be less than pongWait.
+	pingPeriod = (pongWait * 9) / 10
+
+	// Time to wait before force close on connection.
+	closeGracePeriod = 10 * time.Second
+)
+
+var upgrader = websocket.Upgrader{
+	ReadBufferSize:  1024,
+	WriteBufferSize: 1024,
+	// Allow all origins for now - we'll restrict this in production
+	CheckOrigin: func(r *http.Request) bool {
+		return true
+	},
+}
+
+func pumpStdin(ws *websocket.Conn, w io.Writer) {
+	defer ws.Close()
+	ws.SetReadLimit(maxMessageSize)
+	ws.SetReadDeadline(time.Now().Add(pongWait))
+	ws.SetPongHandler(func(string) error { ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
+	for {
+		_, message, err := ws.ReadMessage()
+		if err != nil {
+			break
+		}
+		if _, err := w.Write(message); err != nil {
+			break
+		}
+	}
+}
+
+func pumpStdout(ws *websocket.Conn, r io.Reader, done chan struct{}) {
+	defer func() {
+	}()
+	buf := make([]byte, 1024)
+	for {
+		n, err := r.Read(buf)
+		if err != nil {
+			break
+		}
+		ws.SetWriteDeadline(time.Now().Add(writeWait))
+		if err := ws.WriteMessage(websocket.TextMessage, buf[:n]); err != nil {
+			ws.Close()
+			break
+		}
+	}
+	close(done)
+
+	ws.SetWriteDeadline(time.Now().Add(writeWait))
+	ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
+	time.Sleep(closeGracePeriod)
+	ws.Close()
+}
+
+func ping(ws *websocket.Conn, done chan struct{}) {
+	ticker := time.NewTicker(pingPeriod)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			if err := ws.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(writeWait)); err != nil {
+				log.Println("ping:", err)
+			}
+		case <-done:
+			return
+		}
+	}
+}
+
+func internalError(ws *websocket.Conn, msg string, err error) {
+	log.Println(msg, err)
+	ws.WriteMessage(websocket.TextMessage, []byte("Internal server error."))
+}