├── .gitignore ├── LICENSE ├── README.md ├── cmd ├── app.go ├── clinet.go └── server.go ├── go.mod ├── go.sum ├── main.go └── pkg ├── file ├── mgr.go └── util.go ├── msg ├── msg.go └── pool.go └── server ├── network.go ├── server.go ├── tcp.go └── udp.go /.gitignore: -------------------------------------------------------------------------------- 1 | # Binaries for programs and plugins 2 | *.exe 3 | *.exe~ 4 | *.dll 5 | *.so 6 | *.dylib 7 | 8 | # Test binary, built with `go test -c` 9 | *.test 10 | 11 | # Output of the go coverage tool, specifically when used with LiteIDE 12 | *.out 13 | 14 | # Dependency directories (remove the comment below to include it) 15 | # vendor/ 16 | .idea -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 万彬 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ft 2 | big file transfer, support various network protocols 3 | 4 | ### Use 5 | 6 | ```text 7 | NAME: 8 | ft - big file transfer, support various network protocols 9 | 10 | USAGE: 11 | ft [global options] command [command options] [arguments...] 12 | 13 | COMMANDS: 14 | client, cli start an upload client. 15 | server, srv start a server that receives files and listens on a specified port. 16 | help, h Shows a list of commands or help for one command 17 | 18 | GLOBAL OPTIONS: 19 | --help, -h show help (default: false) 20 | ``` 21 | 22 | ### example 23 | 24 | 启动服务 25 | ```text 26 | ./ft srv|server -d 保存文件的路径 --addr 服务监听的地址 --nw 可选的网络协议 27 | 28 | 29 | NAME: 30 | ft server - start a server that receives files and listens on a specified port. 31 | 32 | USAGE: 33 | ft server [command options] [arguments...] 34 | 35 | OPTIONS: 36 | --addr value specify a listening port (default: "0.0.0.0:9988") 37 | --dir value, -d value upload dir or save dir (default: "./data") 38 | --network value, --nw value choose a network protocol(tcp|udp) (default: "tcp") 39 | ``` 40 | 41 | 42 | 启动客户端 43 | ```text 44 | ./ft cli|client -d 文件所在的文件夹 --addr 服务器地址 --nw 可选的网络协议 [需要传输的文件名,可以多个] 45 | 46 | NAME: 47 | ft client - start an upload client. 48 | 49 | USAGE: 50 | ft client [command options] [arguments...] 51 | 52 | OPTIONS: 53 | --addr value specify a server address (default: "127.0.0.1:9988") 54 | --dir value, -d value upload dir or save dir (default: "./") 55 | --network value, --nw value choose a network protocol(tcp|udp) (default: "tcp") 56 | 57 | ``` 58 | -------------------------------------------------------------------------------- /cmd/app.go: -------------------------------------------------------------------------------- 1 | package cmd 2 | 3 | import ( 4 | "github.com/urfave/cli/v2" 5 | ) 6 | 7 | var commands []*cli.Command 8 | 9 | // 提供其他命令注册 10 | func registerCommand(cmd *cli.Command) { 11 | commands = append(commands, cmd) 12 | } 13 | 14 | // NewApp 创建一个 cli APP，并组装所有的命令。 15 | func NewApp() *cli.App { 16 | app := &cli.App{ 17 | Name: "ft", 18 | Usage: "big file transfer, support various network protocols", 19 | Commands: commands, 20 | EnableBashCompletion: true, 21 | } 22 | 23 | return app 24 | } 25 | -------------------------------------------------------------------------------- /cmd/clinet.go: -------------------------------------------------------------------------------- 1 | package cmd 2 | 3 | import ( 4 | "fmt" 5 | "net" 6 | 7 | "github.com/urfave/cli/v2" 8 | "github.com/ywanbing/ft/pkg/file" 9 | "github.com/ywanbing/ft/pkg/server" 10 | ) 11 | 12 | var clientCmd = &cli.Command{ 13 | Name: "client", 14 | Aliases: []string{"cli"}, 15 | Usage: "start an upload client.", 16 | Flags: []cli.Flag{ 17 | &cli.StringFlag{ 18 | Name: "network", 19 | Aliases: []string{"nw"}, 20 | Usage: "choose a network protocol(tcp|udp)", 21 | Value: "tcp", 22 | }, 23 | &cli.StringFlag{ 24 | Name: "addr", 25 | Usage: "specify a server address", 26 | Value: "127.0.0.1:9988", 27 | }, 28 | &cli.StringFlag{ 29 | Name: "dir", 30 | Value: "./", 31 | Aliases: []string{"d"}, 32 | Usage: "upload dir or save dir", 33 | }, 34 | }, 35 | Action: func(ctx *cli.Context) error { 36 | network := ctx.String("network") 37 | addr := ctx.String("addr") 38 | dir := ctx.String("dir") 39 | 40 | if !file.PathExists(dir) { 41 | return fmt.Errorf("folder does not exist") 42 | } 43 | 44 | fileNames := ctx.Args().Slice() 45 | if len(fileNames) == 0 { 46 | return fmt.Errorf("no transfer files available") 47 | } 48 | 49 | switch network { 50 | case "tcp": 51 | tcpAddr, err := net.ResolveTCPAddr(network, addr) 52 | if err != nil { 53 | return err 54 | } 55 | 56 | conTcp, err := net.DialTCP(network, nil, tcpAddr) 57 | if err != nil { 58 | return err 59 | } 60 | 61 | tcpCon := server.NewTcp(conTcp) 62 | srv := file.NewClient(tcpCon, dir, fileNames) 63 | _ = srv.SendFile() 64 | case "udp": 65 | default: 66 | return fmt.Errorf("network param err: select tcp | udp") 67 | } 68 | return nil 69 | }, 70 | } 71 | 72 | func init() { 73 | registerCommand(clientCmd) 74 | } 75 | -------------------------------------------------------------------------------- /cmd/server.go: -------------------------------------------------------------------------------- 1 | package cmd 2 | 3 | import ( 4 | "fmt" 5 | "net" 6 | "os" 7 | 8 | "github.com/urfave/cli/v2" 9 | "github.com/ywanbing/ft/pkg/file" 10 | "github.com/ywanbing/ft/pkg/server" 11 | ) 12 | 13 | func init() { 14 | registerCommand(serverCmd) 15 | } 16 | 17 | var serverCmd = &cli.Command{ 18 | Name: "server", 19 | Aliases: []string{"srv"}, 20 | Usage: "start a server that receives files and listens on a specified port.", 21 | Flags: []cli.Flag{ 22 | &cli.StringFlag{ 23 | Name: "network", 24 | Aliases: []string{"nw"}, 25 | Usage: "choose a network protocol(tcp|udp)", 26 | Value: "tcp", 27 | }, 28 | &cli.StringFlag{ 29 | Name: "addr", 30 | Usage: "specify a listening port", 31 | Value: "0.0.0.0:9988", 32 | }, 33 | &cli.StringFlag{ 34 | Name: "dir", 35 | Value: "./data", 36 | Aliases: []string{"d"}, 37 | Usage: "upload dir or save dir", 38 | }, 39 | }, 40 | Action: func(ctx *cli.Context) error { 41 | network := ctx.String("network") 42 | addr := ctx.String("addr") 43 | dir := ctx.String("dir") 44 | 45 | if !file.PathExists(dir) { 46 | _ = os.MkdirAll(dir, os.ModePerm) 47 | } 48 | 49 | switch network { 50 | case "tcp": 51 | tcpAddr, err := net.ResolveTCPAddr(network, addr) 52 | if err != nil { 53 | return err 54 | } 55 | 56 | listener, err := net.ListenTCP(network, tcpAddr) 57 | if err != nil { 58 | return err 59 | } 60 | 61 | for { 62 | acceptTCP, err := listener.AcceptTCP() 63 | if err != nil { 64 | return err 65 | } 66 | 67 | fmt.Println("start a connection:", acceptTCP.RemoteAddr()) 68 | tcpCon := server.NewTcp(acceptTCP) 69 | srv := file.NewServer(tcpCon, dir) 70 | _ = srv.Start() 71 | fmt.Println("end a connection:", acceptTCP.RemoteAddr()) 72 | } 73 | case "udp": 74 | default: 75 | return fmt.Errorf("network param err: select tcp | udp") 76 | } 77 | 78 | return nil 79 | }, 80 | } 81 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/ywanbing/ft 2 | 3 | go 1.18 4 | 5 | require ( 6 | github.com/google/uuid v1.2.0 7 | github.com/urfave/cli/v2 v2.15.0 8 | ) 9 | 10 | require ( 11 | github.com/cpuguy83/go-md2man/v2 v2.0.2 // indirect 12 | github.com/russross/blackfriday/v2 v2.1.0 // indirect 13 | github.com/xrash/smetrics v0.0.0-20201216005158-039620a65673 // indirect 14 | ) 15 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/cpuguy83/go-md2man/v2 v2.0.2 h1:p1EgwI/C7NhT0JmVkwCD2ZBK8j4aeHQX2pMHHBfMQ6w= 2 | github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o= 3 | github.com/google/uuid v1.2.0 h1:qJYtXnJRWmpe7m/3XlyhrsLrEURqHRM2kxzoxXqyUDs= 4 | github.com/google/uuid v1.2.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= 5 | github.com/russross/blackfriday/v2 v2.1.0 h1:JIOH55/0cWyOuilr9/qlrm0BSXldqnqwMsf35Ld67mk= 6 | github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= 7 | github.com/urfave/cli/v2 v2.15.0 h1:/U7qTMlBYcmo/Z34PaaVY0Gw04xoGJqEdRAiWNHNyy8= 8 | github.com/urfave/cli/v2 v2.15.0/go.mod h1:1CNUng3PtjQMtRzJO4FMXBQvkGtuYRxxiR9xMa7jMwI= 9 | github.com/xrash/smetrics v0.0.0-20201216005158-039620a65673 h1:bAn7/zixMGCfxrRTfdpNzjtPYqr8smhKouy9mxVdGPU= 10 | github.com/xrash/smetrics v0.0.0-20201216005158-039620a65673/go.mod h1:N3UwUGtsrSj3ccvlPHLoLsHnpR27oXr4ZE984MbSER8= 11 | -------------------------------------------------------------------------------- /main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "log" 5 | "os" 6 | 7 | "github.com/ywanbing/ft/cmd" 8 | ) 9 | 10 | func main() { 11 | if err := cmd.NewApp().Run(os.Args); err != nil { 12 | log.Fatal(err) 13 | } 14 | } 15 | -------------------------------------------------------------------------------- /pkg/file/mgr.go: -------------------------------------------------------------------------------- 1 | package file 2 | 3 | import ( 4 | "fmt" 5 | "io" 6 | "os" 7 | "path/filepath" 8 | "time" 9 | 10 | "github.com/ywanbing/ft/pkg/msg" 11 | "github.com/ywanbing/ft/pkg/server" 12 | ) 13 | 14 | type Server interface { 15 | Start() error 16 | } 17 | 18 | type Client interface { 19 | SendFile() error 20 | } 21 | 22 | type ConMgr struct { 23 | conn server.NetConn 24 | dir string 25 | fileName string 26 | 27 | // 发送的文件列表 28 | sendFiles []string 29 | 30 | // 在发送重要消息的时候，需要同步等待消息的状态，返回是否正确 31 | waitNotify chan bool 32 | stop chan struct{} 33 | } 34 | 35 | func NewServer(conn server.NetConn, dir string) Server { 36 | return &ConMgr{ 37 | conn: conn, 38 | dir: dir, 39 | stop: make(chan struct{}), 40 | } 41 | } 42 | 43 | func (c *ConMgr) Start() error { 44 | c.conn.HandlerLoop() 45 | // 处理接收的消息 46 | return c.handler() 47 | } 48 | 49 | func (c *ConMgr) handler() error { 50 | var fs *os.File 51 | var err error 52 | 53 | defer func() { 54 | if fs != nil { 55 | _ = fs.Close() 56 | } 57 | }() 58 | 59 | for !c.conn.IsClose() { 60 | m, ok := c.conn.GetMsg() 61 | if !ok { 62 | return fmt.Errorf("close by connect") 63 | } 64 | if m == nil { 65 | continue 66 | } 67 | 68 | switch m.MsgType { 69 | case msg.MsgHead: 70 | // 创建文件 71 | if m.FileName != "" { 72 | c.fileName = m.FileName 73 | } else { 74 | c.fileName = GenFileName() 75 | } 76 | 77 | fmt.Println("recv head fileName is", c.fileName) 78 | fs, err = os.OpenFile(filepath.Join(c.dir, c.fileName), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0777) 79 | if err != nil { 80 | fmt.Println("os.Create err =", err) 81 | c.conn.SendMsg(msg.NewNotifyMsg(c.fileName, msg.Status_Err)) 82 | return err 83 | } 84 | fmt.Println("send head is ok") 85 | 86 | c.conn.SendMsg(msg.NewNotifyMsg(c.fileName, msg.Status_Ok)) 87 | case msg.MsgFile: 88 | if fs == nil { 89 | fmt.Println(c.fileName, "file is not open !") 90 | c.conn.SendMsg(msg.NewCloseMsg(c.fileName, msg.Status_Err)) 91 | return nil 92 | } 93 | // 写入文件 94 | _, err = fs.Write(m.Bytes) 95 | if err != nil { 96 | fmt.Println("file.Write err =", err) 97 | c.conn.SendMsg(msg.NewCloseMsg(c.fileName, msg.Status_Err)) 98 | return err 99 | } 100 | case msg.MsgEnd: 101 | // 操作完成 102 | info, _ := fs.Stat() 103 | if info.Size() != int64(m.Size) { 104 | err = fmt.Errorf("file.size %v rece size %v \n", info.Size(), m.Size) 105 | c.conn.SendMsg(msg.NewCloseMsg(c.fileName, msg.Status_Err)) 106 | return err 107 | } 108 | 109 | fmt.Printf("save file %v is success \n", info.Name()) 110 | c.conn.SendMsg(msg.NewNotifyMsg(c.fileName, msg.Status_Ok)) 111 | 112 | fmt.Printf("close file %v is success \n", c.fileName) 113 | _ = fs.Close() 114 | fs = nil 115 | case msg.MsgNotify: 116 | c.waitNotify <- m.Bytes[0] == byte(msg.Status_Ok) 117 | case msg.MsgClose: 118 | fmt.Printf("revc close msg ....\n") 119 | if m.Bytes[0] != byte(msg.Status_Ok) { 120 | return fmt.Errorf("server an error occurred") 121 | } 122 | return nil 123 | } 124 | } 125 | 126 | return err 127 | } 128 | 129 | func NewClient(conn server.NetConn, dir string, files []string) Client { 130 | return &ConMgr{ 131 | conn: conn, 132 | dir: dir, 133 | sendFiles: files, 134 | waitNotify: make(chan bool, 1), 135 | stop: make(chan struct{}), 136 | } 137 | } 138 | 139 | func (c *ConMgr) SendFile() error { 140 | var err error 141 | c.conn.HandlerLoop() 142 | // 处理接收的消息 143 | go func() { 144 | _ = c.handler() 145 | }() 146 | err = c.sendFile() 147 | return err 148 | } 149 | 150 | func (c *ConMgr) sendFile() error { 151 | defer func() { 152 | _ = c.conn.Close() 153 | }() 154 | 155 | var err error 156 | for _, file := range c.sendFiles { 157 | err = c.sendSingleFile(filepath.Join(c.dir, file)) 158 | if err != nil { 159 | return err 160 | } 161 | } 162 | 163 | c.conn.SendMsg(msg.NewCloseMsg(c.fileName, msg.Status_Ok)) 164 | return err 165 | } 166 | 167 | func (c *ConMgr) sendSingleFile(filePath string) error { 168 | file, err := os.Open(filePath) 169 | if err != nil { 170 | fmt.Printf("open file err %v \n", err) 171 | return err 172 | } 173 | 174 | defer func() { 175 | if file != nil { 176 | _ = file.Close() 177 | } 178 | }() 179 | fileInfo, _ := file.Stat() 180 | 181 | fmt.Println("client ready to write ", filePath) 182 | m := msg.NewHeadMsg(fileInfo.Name()) 183 | // 发送文件信息 184 | c.conn.SendMsg(m) 185 | 186 | // 等待服务器返回通知消息 187 | timer := time.NewTimer(5 * time.Second) 188 | select { 189 | case ok := <-c.waitNotify: 190 | if !ok { 191 | return fmt.Errorf("send err") 192 | } 193 | case <-timer.C: 194 | return fmt.Errorf("wait server msg timeout") 195 | } 196 | 197 | for !c.conn.IsClose() { 198 | // 发送文件数据 199 | readBuf := msg.BytesPool.Get().([]byte) 200 | 201 | n, err := file.Read(readBuf) 202 | if err != nil && err != io.EOF { 203 | return err 204 | } 205 | 206 | if n == 0 { 207 | break 208 | } 209 | 210 | c.conn.SendMsg(msg.NewFileMsg(c.fileName, readBuf[:n])) 211 | } 212 | 213 | c.conn.SendMsg(msg.NewEndMsg(c.fileName, uint64(fileInfo.Size()))) 214 | 215 | // 等待服务器返回通知消息 216 | timer = time.NewTimer(5 * time.Second) 217 | select { 218 | case ok := <-c.waitNotify: 219 | if !ok { 220 | return fmt.Errorf("send err") 221 | } 222 | case <-timer.C: 223 | return fmt.Errorf("wait server msg timeout") 224 | } 225 | 226 | fmt.Println("client send " + filePath + " file success...") 227 | return nil 228 | } 229 | -------------------------------------------------------------------------------- /pkg/file/util.go: -------------------------------------------------------------------------------- 1 | package file 2 | 3 | import ( 4 | "os" 5 | 6 | "github.com/google/uuid" 7 | ) 8 | 9 | // PathExists 判断文件夹是否存在 10 | func PathExists(path string) bool { 11 | _, err := os.Stat(path) 12 | if err != nil && os.IsNotExist(err) { 13 | return false 14 | } 15 | return true 16 | } 17 | 18 | func GenFileName() string { 19 | u := uuid.New() 20 | return u.String() 21 | 22 | } 23 | -------------------------------------------------------------------------------- /pkg/msg/msg.go: -------------------------------------------------------------------------------- 1 | package msg 2 | 3 | import ( 4 | "encoding/json" 5 | ) 6 | 7 | type MsgType byte 8 | 9 | const ( 10 | MsgInvalid MsgType = iota 11 | MsgHead 12 | MsgFile 13 | MsgEnd 14 | MsgNotify 15 | MsgClose 16 | ) 17 | 18 | type Status byte 19 | 20 | const ( 21 | Status_Ok Status = iota 22 | Status_Err 23 | ) 24 | 25 | type Message struct { 26 | MsgType MsgType `json:"t"` 27 | FileName string `json:"f"` 28 | Bytes []byte `json:"b"` 29 | Size uint64 `json:"s"` 30 | } 31 | 32 | type Notify struct { 33 | Status byte 34 | } 35 | 36 | func (m *Message) GC() { 37 | if m.MsgType == MsgFile { 38 | BytesPool.Put(m.Bytes[:cap(m.Bytes)]) 39 | } 40 | m.reset() 41 | msgPool.Put(m) 42 | } 43 | 44 | func (m *Message) reset() { 45 | m.MsgType = MsgInvalid 46 | m.FileName = "" 47 | m.Bytes = nil 48 | m.Size = 0 49 | } 50 | 51 | func (m *Message) String() string { 52 | bytes, _ := json.Marshal(m) 53 | return string(bytes) 54 | } 55 | 56 | // Decode will convert from bytes 57 | func Decode(b []byte) (m *Message, err error) { 58 | m = msgPool.Get().(*Message) 59 | err = json.Unmarshal(b, &m) 60 | return 61 | } 62 | 63 | func NewNotifyMsg(fileName string, status Status) *Message { 64 | m := msgPool.Get().(*Message) 65 | m.MsgType = MsgNotify 66 | m.Bytes = []byte{byte(status)} 67 | m.FileName = fileName 68 | return m 69 | } 70 | 71 | func NewHeadMsg(fileName string) *Message { 72 | m := msgPool.Get().(*Message) 73 | m.MsgType = MsgHead 74 | m.FileName = fileName 75 | return m 76 | } 77 | 78 | func NewFileMsg(fileName string, buf []byte) *Message { 79 | m := msgPool.Get().(*Message) 80 | m.MsgType = MsgFile 81 | m.FileName = fileName 82 | m.Bytes = buf 83 | // m.Bytes = make([]byte, len(buf)) 84 | // copy(m.Bytes, buf) 85 | return m 86 | } 87 | 88 | func NewEndMsg(fileName string, size uint64) *Message { 89 | m := msgPool.Get().(*Message) 90 | m.MsgType = MsgEnd 91 | m.FileName = fileName 92 | m.Size = size 93 | return m 94 | } 95 | 96 | func NewCloseMsg(fileName string, status Status) *Message { 97 | m := msgPool.Get().(*Message) 98 | m.MsgType = MsgClose 99 | m.Bytes = []byte{byte(status)} 100 | m.FileName = fileName 101 | return m 102 | } 103 | -------------------------------------------------------------------------------- /pkg/msg/pool.go: -------------------------------------------------------------------------------- 1 | package msg 2 | 3 | import ( 4 | "sync" 5 | ) 6 | 7 | var ( 8 | msgPool = sync.Pool{ 9 | New: func() any { 10 | return &Message{} 11 | }, 12 | } 13 | 14 | BytesPool = sync.Pool{ 15 | New: func() any { 16 | return make([]byte, 40*1024) 17 | }, 18 | } 19 | ) 20 | -------------------------------------------------------------------------------- /pkg/server/network.go: -------------------------------------------------------------------------------- 1 | package server 2 | 3 | import ( 4 | "github.com/ywanbing/ft/pkg/msg" 5 | ) 6 | 7 | type NetConn interface { 8 | // HandlerLoop 不能阻塞 9 | HandlerLoop() 10 | GetMsg() (*msg.Message, bool) 11 | SendMsg(m *msg.Message) 12 | Close() error 13 | IsClose() bool 14 | } 15 | -------------------------------------------------------------------------------- /pkg/server/server.go: -------------------------------------------------------------------------------- 1 | package server 2 | 3 | import ( 4 | "fmt" 5 | ) 6 | 7 | var ( 8 | MAGIC_BYTES = []byte("f00t") 9 | EmErr = fmt.Errorf("dont have msg") 10 | ) 11 | -------------------------------------------------------------------------------- /pkg/server/tcp.go: -------------------------------------------------------------------------------- 1 | package server 2 | 3 | import ( 4 | "bytes" 5 | "encoding/binary" 6 | "fmt" 7 | "io" 8 | "net" 9 | "sync" 10 | "time" 11 | 12 | "github.com/ywanbing/ft/pkg/msg" 13 | ) 14 | 15 | type TcpCon struct { 16 | conn *net.TCPConn 17 | 18 | recv chan *msg.Message 19 | send chan *msg.Message 20 | 21 | onceStop *sync.Once 22 | stop chan struct{} 23 | } 24 | 25 | func NewTcp(conn *net.TCPConn) *TcpCon { 26 | return &TcpCon{ 27 | conn: conn, 28 | recv: make(chan *msg.Message, 1024), 29 | send: make(chan *msg.Message, 1024), 30 | onceStop: &sync.Once{}, 31 | stop: make(chan struct{}), 32 | } 33 | } 34 | 35 | func (t *TcpCon) HandlerLoop() { 36 | go t.readMsg() 37 | go t.sendMsg() 38 | } 39 | 40 | func (t *TcpCon) sendMsg() { 41 | var err error 42 | defer func() { 43 | if err != nil { 44 | fmt.Printf("found mistake: %s \n", err) 45 | } 46 | _ = t.Close() 47 | }() 48 | 49 | buf := make([]byte, 64*1024) 50 | ticker := time.NewTicker(5 * time.Second) 51 | defer ticker.Stop() 52 | 53 | for !t.IsClose() { 54 | select { 55 | case m := <-t.send: 56 | data := m.String() 57 | m.GC() 58 | 59 | dataLen := len(data) 60 | 61 | copy(buf[:4], MAGIC_BYTES) 62 | binary.BigEndian.PutUint32(buf[4:8], uint32(dataLen)) 63 | copy(buf[8:], []byte(data)) 64 | 65 | _, err = t.conn.Write(buf[:8+dataLen]) 66 | if err != nil { 67 | return 68 | } 69 | case <-ticker.C: 70 | 71 | } 72 | } 73 | } 74 | 75 | func (t *TcpCon) readMsg() { 76 | var err error 77 | defer func() { 78 | if err != nil { 79 | fmt.Printf("found mistake: %s \n", err) 80 | } 81 | _ = t.Close() 82 | }() 83 | 84 | header := make([]byte, 4) 85 | buf := make([]byte, 64*1024) 86 | 87 | for { 88 | // read until we get 4 bytes for the magic 89 | _, err = io.ReadFull(t.conn, header) 90 | if err != nil { 91 | if err != io.EOF { 92 | err = fmt.Errorf("initial read error: %v \n", err) 93 | return 94 | } 95 | time.Sleep(10 * time.Millisecond) 96 | continue 97 | } 98 | 99 | if !bytes.Equal(header, MAGIC_BYTES) { 100 | err = fmt.Errorf("initial bytes are not magic: %s", header) 101 | return 102 | } 103 | 104 | // read until we get 4 bytes for the header 105 | _, err = io.ReadFull(t.conn, header) 106 | if err != nil { 107 | err = fmt.Errorf("initial read error: %v \n", err) 108 | return 109 | } 110 | 111 | // 数据大小 112 | msgSize := binary.BigEndian.Uint32(header) 113 | 114 | // 解析为结构体消息 115 | var n int 116 | var m *msg.Message 117 | 118 | n, err = io.ReadFull(t.conn, buf[:msgSize]) 119 | if err != nil { 120 | err = fmt.Errorf("initial read error: %v \n", err) 121 | return 122 | } 123 | 124 | m, err = msg.Decode(buf[:n]) 125 | if err != nil { 126 | err = fmt.Errorf("read message error: %v \n", err) 127 | return 128 | } 129 | 130 | t.recv <- m 131 | } 132 | } 133 | 134 | func (t *TcpCon) GetMsg() (*msg.Message, bool) { 135 | timer := time.NewTimer(5 * time.Second) 136 | defer timer.Stop() 137 | select { 138 | case m, ok := <-t.recv: 139 | return m, ok 140 | case <-timer.C: 141 | return nil, true 142 | } 143 | } 144 | 145 | func (t *TcpCon) SendMsg(m *msg.Message) { 146 | t.send <- m 147 | } 148 | 149 | func (t *TcpCon) Close() error { 150 | t.onceStop.Do(func() { 151 | fmt.Println("close a connect, addr: ", t.conn.RemoteAddr()) 152 | _ = t.conn.Close() 153 | close(t.stop) 154 | }) 155 | return nil 156 | } 157 | 158 | func (t *TcpCon) IsClose() bool { 159 | select { 160 | case <-t.stop: 161 | return true 162 | default: 163 | return false 164 | } 165 | } 166 | 167 | var _ = NetConn(&TcpCon{}) 168 | -------------------------------------------------------------------------------- /pkg/server/udp.go: -------------------------------------------------------------------------------- 1 | package server 2 | --------------------------------------------------------------------------------