websocket转socket

websocket协议转tcp协议的代理
背景：
原先有个页游的项目，需要移植到手机上，做手机页游，也就是到html5（h5)版本，面临一个问题：原先服务器和客户端（flash）通信用的是tcp协议，而h5用的是websocket，虽然也也是用来传输二进制的，但毕竟原先的服务器不认识websocket的握手协议以及中间传输的帧协议，又不想改动服务器的c++代码来做这个恶心的适配，怎么办？
解决：
组内最近有同事在研究go，就让其用go语言实现了一个websocket代理，把websocket数据包去掉websocket的协议内容，只把真正的数据转发到原先的服务器上。
代码很简单，经过测试性能也不错，代码丢在了github上，有需要的可以去看： https://github.com/zxb19890727/websocket_transfer#websocket_transfer

package main

import (
    "io"
    "log"
    "net"
    "os"
    "fmt"
    "bufio"
    "bytes"
    "strconv"
    "strings"
    "crypto/sha1"
    "runtime"
    "encoding/base64"
    "encoding/binary"
    "net/http"
)

type WebSocket struct {
    Listener    net.Listener
    Clients        []*Client
}

var num int
var new_addr string
type Client struct {
    Conn         net.Conn
    Nickname    string
    Shook        bool
    Server        *WebSocket
    Id            int
    TcpConn      net.Conn
    WebsocketType   int
    Num             int
}

type Msg struct {
    Data        string
    Num            int
}

func (self *Client) Release() {
    // release all connect
    self.TcpConn.Close()
    self.Conn.Close()
}


func (self *Client) Handle() {
    defer self.Release()
    if !self.Handshake() {
        // handshak err , del this conn
        return
    }

    // connect to another server for tcp
    if !self.ConnTcpServer(){
        // can not connect to the other server , release
        return
    }
    num = num + 1
    log.Print("now connect num : ", num)
    self.Num = num
    go self.Read()
    self.ReadTcp()
}

func (self *Client) Read() {
    var (
        buf     []byte
        err     error
        rsv     byte
        opcode    byte
        mask    byte
        mKey    []byte
        length    uint64
        l        uint16
        payload    byte
    )
    for {
        buf = make([]byte, 2)
        _, err = io.ReadFull(self.Conn, buf)
        if err != nil {
            self.Release()
            break
        }
        //fin = buf[0] >> 7
        //if fin == 0 {
        //}

        rsv = (buf[0] >>4) &0x7
        // which must be 0
        if rsv != 0{
            log.Print("Client send err msg:",rsv,", disconnect it")
            self.Release()
            break
        }

        opcode = buf[0] & 0xf
        // opcode   if 8 then disconnect
        if opcode == 8 {
            log.Print("CLient want close Connection")
            self.Release()
            break
        }

        // should save the opcode 
        // if client send by binary should return binary (especially for Egret)
        self.WebsocketType = int(opcode)

        mask = buf[1] >> 7
        // the translate may have mask 

        payload = buf[1] & 0x7f
        // if length < 126 then payload mean the length
        // if length == 126 then the next 8bit mean the length
        // if length == 127 then the next 64bit mean the length
        switch {
        case payload < 126:
            length = uint64(payload)

        case payload == 126:
            buf = make([]byte, 2)
            io.ReadFull(self.Conn, buf)
            binary.Read(bytes.NewReader(buf), binary.BigEndian, &l)
            length = uint64(l)

        case payload == 127:
            buf = make([]byte, 8)
            io.ReadFull(self.Conn, buf)
            binary.Read(bytes.NewReader(buf), binary.BigEndian, &length)
        }
        if mask == 1 {
            mKey = make([]byte, 4)
            io.ReadFull(self.Conn, mKey)
        }
        buf = make([]byte, length)
        io.ReadFull(self.Conn, buf)
        if mask == 1 {
            for i, v := range buf {
                buf[i] = v ^ mKey[i % 4]
            }
            //fmt.Print("mask", mKey)
        }
        log.Print("rec from the client(",self.Num,")", string(buf))
        self.TcpConn.Write(buf)
    }
}
// read from other tcp
func (self *Client) ReadTcp() {
    var (
        buf  []byte
    )
    buf = make([]byte, 1024)

    for {
        length,err := self.TcpConn.Read(buf)

        if err != nil {
            self.Release()
            num = num - 1
            // only 
            log.Print("other tcp connect err", err)
            log.Print("disconnect client :", self.Num)
            log.Print("now have:", num)
            break
        }
        log.Print("recv from other tcp : ", string(buf[:length]))
        self.Write(buf[:length])
        //Write to websocket
    }
}

// write to websocket
func (self *Client) Write(data []byte) bool {
    data_binary := new(bytes.Buffer) //which 

    //should be binary or string
    frame := []byte{129}  //string
    length := len(data)
    // 10000001
    if self.WebsocketType == 2 {
        frame = []byte{130}
        // 10000010
        err := binary.Write(data_binary, binary.LittleEndian, data)
        if err != nil {
            log.Print(" translate to binary err", err)
        }
        length = len(data_binary.Bytes())
    }
    switch {
    case length < 126:
        frame = append(frame, byte(length))
    case length <= 0xffff:
        buf := make([]byte, 2)
        binary.BigEndian.PutUint16(buf, uint16(length))
        frame = append(frame, byte(126))
        frame = append(frame, buf...)
    case uint64(length) <= 0xffffffffffffffff:
        buf := make([]byte, 8)
        binary.BigEndian.PutUint64(buf, uint64(length))
        frame = append(frame, byte(127))
        frame = append(frame, buf...)
    default:
        log.Print("Data too large")
        return false
    }
    if self.WebsocketType == 2 {
        frame = append(frame, data_binary.Bytes()...)
    } else {
        frame = append(frame, data...)
    }
    self.Conn.Write(frame)
    frame = []byte{0}
    return true
}

func (self *Client) ConnTcpServer() bool {

    conn, err := net.Dial("tcp", new_addr)

    if(err != nil) {
        log.Print("connect other tcp server error")
        return false
    }

    self.TcpConn = conn
    return true
}

func (self *Client) Handshake() bool {
    if self.Shook {
        return true
    }
    reader := bufio.NewReader(self.Conn)
    key := ""
    str := ""
    for {
        line, _, err := reader.ReadLine()
        if err != nil {
            log.Print("Handshake err:",err)
            return false
        }
        if len(line) == 0 {
            break
        }
        str = string(line)
        if strings.HasPrefix(str, "Sec-WebSocket-Key") {
            if len(line)>= 43 {
                key = str[19:43]
            }
        }
    }
    if key == "" {
        return false
    }
    sha := sha1.New()
    io.WriteString(sha, key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11")
    key = base64.StdEncoding.EncodeToString(sha.Sum(nil))
    header := "HTTP/1.1 101 Switching Protocols\r\n" +
    "Connection: Upgrade\r\n" +
    "Sec-WebSocket-Version: 13\r\n" +
    "Sec-WebSocket-Accept: " + key + "\r\n" +
    "Upgrade: websocket\r\n\r\n"
    self.Conn.Write([]byte(header))
    self.Shook = true
    self.Server.Clients = append(self.Server.Clients, self)
    return true
}

func NewWebSocket(addr string) *WebSocket {
    l, err := net.Listen("tcp", addr)
    if err != nil {
        log.Fatal(err)
        // if cannot listen then return
    }
    return &WebSocket{l, make([]*Client, 0)}
}


func (self *WebSocket) Loop() {
    for {
        conn, err := self.Listener.Accept()
        if err != nil {
            log.Print("client conn err:", err)
            continue
        }
        s := conn.RemoteAddr().String()
        i, _ := strconv.Atoi(strings.Split(s, ":")[1])
        client := &Client{conn, "", false, self, i, conn, 1, num}
        go client.Handle()
    }
}


func handler(w http.ResponseWriter, r *http.Request) {
    // show num of goroutine
    w.Header().Set("Content-Type", "text/plain")
    num := strconv.FormatInt(int64(runtime.NumGoroutine()), 10)
    w.Write([]byte(num))
}

func main() {
    arg_num:=len(os.Args)
    if arg_num < 2 {
        fmt.Println(arg_num)
        fmt.Print("Wrong Arguments\nxxxx xxx.xxx.xxx.xxx:xxxx\nport ip:port(for tcp)")
        os.Exit(0)
    }
    num = 0
    conn,err := net.Dial("tcp", string(os.Args[2]))
    if err != nil {
        fmt.Println(err)
        os.Exit(0)
    }
    new_addr = os.Args[2]
    fmt.Println("Check Arguments Ok")
    conn.Close()
    port := os.Args[1]
    ip_port:= "0.0.0.0:"+string(port)
    ws := NewWebSocket(ip_port)
    // listen 9051
    go  ws.Loop()
    fmt.Println("Start Listen")
    // listen 11181 to show num of goroutine
    http.HandleFunc("/", handler)
    http.ListenAndServe(":11181", nil)
}
最后编辑于 ：
©著作权归作者所有,转载或内容合作请联系作者
  • 人面猴
    序言：七十年代末，一起剥皮案震惊了整个滨河市，随后出现的几起案子，更是在滨河造成了极大的恐慌，老刑警刘岩，带你破解...
    沈念sama阅读 199,340评论 5赞 467
  • 死咒
    序言：滨河连续发生了三起死亡事件，死亡现场离奇诡异，居然都是意外死亡，警方通过查阅死者的电脑和手机，发现死者居然都...
    沈念sama阅读 83,762评论 2赞 376
  • 救了他两次的神仙让他今天三更去死
    文/潘晓璐 我一进店门，熙熙楼的掌柜王于贵愁眉苦脸地迎上来，“玉大人，你说我怎么就摊上这事。” “怎么了？”我有些...
    开封第一讲书人阅读 146,329评论 0赞 329
  • 道士缉凶录：失踪的卖姜人
    文/不坏的土叔 我叫张陵，是天一观的道长。 经常有香客问我，道长，这世上最难降的妖魔是什么？ 我笑而不...
    开封第一讲书人阅读 53,678评论 1赞 270
  • 港岛之恋（遗憾婚礼）
    正文 为了忘掉前任，我火速办了婚礼，结果婚礼上，老公的妹妹穿的比我还像新娘。我一直安慰自己，他们只是感情好，可当我...
    茶点故事阅读 62,583评论 5赞 359
  • 恶毒庶女顶嫁案：这布局不是一般人想出来的
    文/花漫 我一把揭开白布。 她就那样静静地躺着，像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上，一...
    开封第一讲书人阅读 47,995评论 1赞 275
  • 城市分裂传说
    那天，我揣着相机与录音，去河边找鬼。 笑死，一个胖子当着我的面吹牛，可吹牛的内容都是我干的。 我是一名探鬼主播，决...
    沈念sama阅读 37,493评论 3赞 390
  • 双鸳鸯连环套：你想象不到人心有多黑
    文/苍兰香墨 我猛地睁开眼，长吁一口气：“原来是场噩梦啊……” “哼！你这毒妇竟也来了？” 一声冷哼从身侧响起，我...
    开封第一讲书人阅读 36,145评论 0赞 254
  • 万荣杀人案实录
    序言：老挝万荣一对情侣失踪，失踪者是张志新（化名）和其女友刘颖，没想到半个月后，有当地人在树林里发现了一具尸体，经...
    沈念sama阅读 40,293评论 1赞 294
  • 护林员之死
    正文 独居荒郊野岭守林人离奇死亡，尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,250评论 2赞 317
  • 白月光启示录
    正文 我和宋清朗相恋三年，在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,267评论 1赞 328
  • 活死人
    序言：一个原本活蹦乱跳的男人离奇死亡，死状恐怖，灵堂内的尸体忽然破棺而出，到底是诈尸还是另有隐情，我是刑警宁泽，带...
    沈念sama阅读 32,973评论 3赞 316
  • 日本核电站爆炸内幕
    正文 年R本政府宣布，位于F岛的核电站，受9级特大地震影响，放射性物质发生泄漏。R本人自食恶果不足惜，却给世界环境...
    茶点故事阅读 38,556评论 3赞 303
  • 男人毒药：我在死后第九天来索命
    文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹，春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,648评论 0赞 19
  • 一桩弑父案，背后竟有这般阴谋
    文/苍兰香墨 我抬头看了看天上的太阳。三九已至，却和暖如春，着一层夹袄步出监牢的瞬间，已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 30,873评论 1赞 255
  • 情欲美人皮
    我被黑心中介骗来泰国打工， 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留，地道东北人。 一个月前我还...
    沈念sama阅读 42,257评论 2赞 345
  • 代替公主和亲
    正文 我出身青楼，却偏偏与公主长得像，于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子，可洞房花烛夜当晚...
    茶点故事阅读 41,809评论 2赞 339