ws.go

go.patch

diff -r d5666bad617d src/pkg/net/file.go
--- a/src/pkg/net/file.go	Thu Sep 27 15:36:22 2012 -0700
+++ b/src/pkg/net/file.go	Wed Apr 24 11:49:35 2013 -0400
@@ -7,10 +7,132 @@
 package net
 
 import (
+	"errors"
+	"io"
 	"os"
 	"syscall"
+	"time"
 )
 
+type ptyAddr int
+
+func (ptyAddr) Network() string {
+	return "pty"
+}
+
+func (ptyAddr) String() string {
+	return "pty"
+}
+
+type PtyConn struct {
+	fd *netFD
+}
+
+func newPtyConn(fd *netFD) *PtyConn {
+	return &PtyConn{fd}
+}
+
+func (c *PtyConn) ok() bool { return c != nil && c.fd != nil }
+
+// Implementation of the Conn interface - see Conn for documentation.
+
+// Read implements the Conn Read method.
+func (c *PtyConn) Read(b []byte) (n int, err error) {
+	if !c.ok() {
+		return 0, syscall.EINVAL
+	}
+	return c.fd.Read(b)
+}
+
+// ReadFrom implements the io.ReaderFrom ReadFrom method.
+func (c *PtyConn) ReadFrom(r io.Reader) (int64, error) {
+	if n, err, handled := sendFile(c.fd, r); handled {
+		return n, err
+	}
+	return genericReadFrom(c, r)
+}
+
+// Write implements the Conn Write method.
+func (c *PtyConn) Write(b []byte) (n int, err error) {
+	if !c.ok() {
+		return 0, syscall.EINVAL
+	}
+	return c.fd.Write(b)
+}
+
+// Close closes the TCP connection.
+func (c *PtyConn) Close() error {
+	if !c.ok() {
+		return syscall.EINVAL
+	}
+	return c.fd.Close()
+}
+
+// LocalAddr returns the local network address, a *TCPAddr.
+func (c *PtyConn) LocalAddr() Addr {
+	return ptyAddr(0)
+}
+
+// RemoteAddr returns the remote network address, a *TCPAddr.
+func (c *PtyConn) RemoteAddr() Addr {
+	return ptyAddr(0)
+}
+
+// SetDeadline implements the Conn SetDeadline method.
+func (c *PtyConn) SetDeadline(t time.Time) error {
+	return errors.New("net.Pipe does not support deadlines")
+}
+
+// SetReadDeadline implements the Conn SetReadDeadline method.
+func (c *PtyConn) SetReadDeadline(t time.Time) error {
+	return errors.New("net.Pipe does not support deadlines")
+}
+
+// SetWriteDeadline implements the Conn SetWriteDeadline method.
+func (c *PtyConn) SetWriteDeadline(t time.Time) error {
+	return errors.New("net.Pipe does not support deadlines")
+}
+
+// File returns a copy of the underlying os.File, set to blocking mode.
+// It is the caller's responsibility to close f when finished.
+// Closing c does not affect f, and closing f does not affect c.
+func (c *PtyConn) File() (f *os.File, err error) { return c.fd.dup() }
+
+func newPtyFD(f *os.File) (*netFD, error) {
+	syscall.ForkLock.RLock()
+	fd, err := syscall.Dup(int(f.Fd()))
+	if err != nil {
+		syscall.ForkLock.RUnlock()
+		return nil, os.NewSyscallError("dup", err)
+	}
+	syscall.CloseOnExec(fd)
+	syscall.ForkLock.RUnlock()
+
+	// We want blocking mode for the new fd, hence the double negative.
+	if err = syscall.SetNonblock(fd, false); err != nil {
+		closesocket(fd)
+		return nil, err
+	}
+
+	family := syscall.AF_UNSPEC
+
+	netfd, err := newFD(fd, family, syscall.SOCK_SEQPACKET, "pty")
+	if err != nil {
+		closesocket(fd)
+		return nil, err
+	}
+	netfd.setAddr(ptyAddr(0), ptyAddr(0))
+	return netfd, nil
+}
+
+func NewPtyConn(f *os.File) (c *PtyConn, err error) {
+	fd, err := newPtyFD(f)
+	if err != nil {
+		return nil, err
+	}
+	return newPtyConn(fd), nil
+}
+
 func newFileFD(f *os.File) (*netFD, error) {
 	syscall.ForkLock.RLock()
 	fd, err := syscall.Dup(int(f.Fd()))

ws.go

package main

import (
	"code.google.com/p/go.net/websocket"
	"encoding/json"
	"github.com/kghost/pty"
	"log"
	"net"
	"net/http"
	"os/exec"
	"sync"
	"syscall"
	"unsafe"
)

func marshal(v interface{}) (msg []byte, payloadType byte, err error) {
	switch data := v.(type) {
	case []byte:
		return data, websocket.BinaryFrame, nil
	}
	return nil, websocket.UnknownFrame, websocket.ErrNotSupported
}

func unmarshal(msg []byte, payloadType byte, v interface{}) (err error) {
	switch data := v.(type) {
	case func(msg []byte, payloadType byte) (err error):
		return data(msg, payloadType)
	}
	return websocket.ErrNotSupported
}

var message = websocket.Codec{marshal, unmarshal}

func echoServer(ws *websocket.Conn) {
	log.Printf("connected: %s\n", ws.RemoteAddr())
	c := exec.Command("/bin/login", "-f", "kghost")
	fd, err := pty.Start(c)
	if err != nil {
		log.Printf("pty error: %s\n", err)
		return
	}

	master, err := net.NewPtyConn(fd)
	if err != nil {
		log.Printf("pty error: %s\n", err)
		fd.Close()
		return
	}

	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		var err error = nil
		for err == nil {
			err = message.Receive(ws, func(msg []byte, payloadType byte) (err error) {
				switch payloadType {
				case websocket.BinaryFrame:
					_, ew := master.Write(msg)
					return ew
				case websocket.TextFrame:
					m := map[string]interface{}{}
					json.Unmarshal(msg, &m)
					cmd, ok1 := m["cmd"].(float64)
					if ok1 {
						switch int(cmd) {
						case 0:
							w, ok2 := m["width"].(float64)
							h, ok3 := m["height"].(float64)
							if ok2 && ok3 {
								type winsize struct {
									ws_row, ws_col       uint16
									ws_xpixel, ws_ypixel uint16
								}

								ws := winsize{
									ws_row: uint16(h),
									ws_col: uint16(w),
								}
								syscall.Syscall(syscall.SYS_IOCTL,
									fd.Fd(), uintptr(syscall.TIOCSWINSZ),
									uintptr(unsafe.Pointer(&ws)))
								return nil
							} else {
								return websocket.ErrNotSupported
							}
						}
						return websocket.ErrNotSupported
					} else {
						return websocket.ErrNotSupported
					}
				}
				return websocket.ErrNotSupported
			})
		}
		log.Printf("ws -> sh: %s\n", err)
		master.Close()
		fd.Close()
		ws.Close()
		wg.Done()
	}()
	go func() {
		var err error = nil
		for err == nil {
			var c int = 0
			var bs []byte = make([]byte, 1024, 1024)
			c, err = master.Read(bs)
			if err != nil {
				break
			}
			if c != 0 {
				err = message.Send(ws, bs[0:c])
				if err != nil {
					break
				}
			}
		}
		log.Printf("sh -> ws: %s\n", err)
		master.Close()
		fd.Close()
		ws.Close()
		wg.Done()
	}()
	wg.Wait()
	log.Printf("disconnected: %s\n", ws.RemoteAddr())
	c.Wait()
}

func main() {
	http.Handle("/", websocket.Handler(echoServer))
	err := http.ListenAndServe("localhost:12345", nil)
	if err != nil {
		panic("ListenAndServe: " + err.Error())
	}
}