TCP/IP Networking in Go

package main


import (
  "bufio"
  "io"
  "log"
  "net"
  "strconv"
  "strings"
  "sync"


  "github.com/pkg/errors"


  "encoding/gob"
  "flag"
)
复制

type complexData struct {
  N int
  S string
  M map[string]int
  P []byte
  C *complexData
}
复制

const (
    Port = ":61000"
)
复制

func Open(addr string) (*bufio.ReadWriter, error) {
  log.Println("Dial " + addr)
  conn, err := net.Dial("tcp", addr)
  if err != nil {
    return nil, errors.Wrap(err, "Dialing "+addr+" failed")
  }
  return bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn)), nil
}
复制

type HandleFunc func(*bufio.ReadWriter)
复制

type Endpoint struct {
  listener net.Listener
  handler  map[string]HandleFunc
    // map不是线程安全的，需要一个mutex来控制访问
  m sync.RWMutex
}
复制

func NewEndpoint() *Endpoint {
  return &Endpoint{
    handler: map[string]HandleFunc{},
  }
}
复制

func (e *Endpoint) AddHandleFunc(name string, f HandleFunc) {
  e.m.Lock()
  e.handler[name] = f
  e.m.Unlock()
}
复制

func (e *Endpoint) Listen() error {
  var err error
  e.listener, err = net.Listen("tcp", Port)
  if err != nil {
    return errors.Wrapf(err, "Unable to listen on port %s\n", Port)
  }
  log.Println("Listen on", e.listener.Addr().String())
  for {
    log.Println("Accept a connection request.")
    conn, err := e.listener.Accept()
    if err != nil {
      log.Println("Failed accepting a connection request:", err)
      continue
    }
    log.Println("Handle incoming messages.")
    go e.handleMessages(conn)
  }
}
复制

func (e *Endpoint) handleMessages(conn net.Conn) {
  // 使用buffered reader包裹，以方便读取数据
  rw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
  defer conn.Close()
  // 解析出command,执行对应的处理器函数
  for {
  log.Print("Receive command '")
  cmd, err := rw.ReadString('\n')
  switch {
  case err == io.EOF:
    log.Println("Reached EOF - close this connection.\n   ---")
    return
  case err != nil:
    log.Println("\nError reading command. Got: '"+cmd+"'\n", err)
    return
  }
  //移除行尾的换行符
  cmd = strings.Trim(cmd, "\n ")
  log.Println(cmd + "'")
    
  // 找到command对应的处理器函数，并调用它
  e.m.RLock()
  handleCommand, ok := e.handler[cmd]
  e.m.RUnlock()
  if !ok {
    log.Println("Command '" + cmd + "' is not registered.")
    return
  }
  handleCommand(rw)
  }
}
复制

func handleStrings(rw *bufio.ReadWriter) {
  log.Print("Receive STRING message:")
  s, err := rw.ReadString('\n')
  if err != nil {
    log.Println("Cannot read from connection.\n", err)
  }
  s = strings.Trim(s, "\n ")
  log.Println(s)
  _, err = rw.WriteString("Thank you.\n")
  if err != nil {
    log.Println("Cannot write to connection.\n", err)
  }
  err = rw.Flush()
  if err != nil {
    log.Println("Flush failed.", err)
  }
}
复制

func handleGob(rw *bufio.ReadWriter) {
  log.Print("Receive GOB data:")
  var data complexData
  // 创建一个解码器，直接解码为一个struct变量
  dec := gob.NewDecoder(rw)
  err := dec.Decode(&data)
  if err != nil {
    log.Println("Error decoding GOB data:", err)
    return
  }
  // 打印complexData以及深一层的数据，证明数据确实传输过来了
  log.Printf("Outer complexData struct: \n%#v\n", data)


  log.Printf("Inner complexData struct: \n%#v\n", data.C)
}
复制

func client(ip string) error {
  // 测试数据，注意GOB能正确处理map、slice、struct等数据类型
  testStruct := complexData{
    N: 23,
    S: "string data",
    M: map[string]int{"one": 1, "two": 2, "three": 3},
    P: []byte("abc"),
    C: &complexData{
      N: 256,
      S: "Recursive structs? Piece of cake!",
      M: map[string]int{"01": 1, "10": 2, "11": 3},
    },
  }
    // 连接到服务端
  rw, err := Open(ip + Port)
  if err != nil {
    return errors.Wrap(err, "Client: Failed to open connection to "+ip+Port)
  }
    // 发送 STRING 请求。先发送命令，再发送数据
  log.Println("Send the string request.")
  n, err := rw.WriteString("STRING\n")
  if err != nil {
    return errors.Wrap(err, "Could not send the STRING request ("+strconv.Itoa(n)+" bytes written)")
  }
  n, err = rw.WriteString("Additional data.\n")
  if err != nil {
    return errors.Wrap(err, "Could not send additional STRING data ("+strconv.Itoa(n)+" bytes written)")
  }
  log.Println("Flush the buffer.")
  err = rw.Flush()
  if err != nil {
    return errors.Wrap(err, "Flush failed.")
  }
  // 读取响应数据
  log.Println("Read the reply.")
  response, err := rw.ReadString('\n')
  if err != nil {
    return errors.Wrap(err, "Client: Failed to read the reply: '"+response+"'")
  }


  log.Println("STRING request: got a response:", response)
    
  // 发送 GOB 请求。创建一个编码器，直接编码到rw。发送 命令名 ，然后发送 GOB数据
  log.Println("Send a struct as GOB:")
  log.Printf("Outer complexData struct: \n%#v\n", testStruct)
  log.Printf("Inner complexData struct: \n%#v\n", testStruct.C)
  enc := gob.NewEncoder(rw)
  n, err = rw.WriteString("GOB\n")
  if err != nil {
    return errors.Wrap(err, "Could not write GOB data ("+strconv.Itoa(n)+" bytes written)")
  }
  err = enc.Encode(testStruct)
  if err != nil {
    return errors.Wrapf(err, "Encode failed for struct: %#v", testStruct)
  }
  err = rw.Flush()
  if err != nil {
    return errors.Wrap(err, "Flush failed.")
  }
  return nil
}
复制

func server() error {
  endpoint := NewEndpoint()
   注册处理器函数
  endpoint.AddHandleFunc("STRING", handleStrings)
  endpoint.AddHandleFunc("GOB", handleGob)
   开始监听
  return endpoint.Listen()
}
复制

func main() {
  connect := flag.String("connect", "", "IP address of process to join. If empty, go into listen mode.")
  flag.Parse()
  // 如果connect选项不为空，进入客户端模式
  if *connect != "" {
    err := client(*connect)
    if err != nil {
      log.Println("Error:", errors.WithStack(err))
    }
    log.Println("Client done.")
    return
  }
    // 否则进入服务端模式
  err := server()
  if err != nil {
    log.Println("Error:", errors.WithStack(err))
  }


  log.Println("Server done.")
}
复制

func init() {
  log.SetFlags(log.Lshortfile)
}
复制