From 3f7b57740d2873f1bcebecd6b3d11e4edceec928 Mon Sep 17 00:00:00 2001 From: "arraykeys@gmail.com" Date: Mon, 25 Sep 2017 15:31:07 +0800 Subject: [PATCH] Signed-off-by: arraykeys@gmail.com --- config.go | 23 +++--- services/args.go | 16 ++-- services/http.go | 63 +++++++------- services/service.go | 6 +- services/tcp.go | 103 +++++++++++++++++------ services/tls.go | 117 -------------------------- services/udp.go | 194 +++++++++++++++++++++++++++++++++++++++++++- utils/functions.go | 46 +++++++++++ 8 files changed, 370 insertions(+), 198 deletions(-) delete mode 100644 services/tls.go diff --git a/config.go b/config.go index 1f18c94..4fd2856 100755 --- a/config.go +++ b/config.go @@ -13,7 +13,7 @@ import ( var ( app *kingpin.Application - service services.ServiceItem + service *services.ServiceItem ) func initConfig() (err error) { @@ -21,7 +21,7 @@ func initConfig() (err error) { //define args tcpArgs := services.TCPArgs{} httpArgs := services.HTTPArgs{} - tlsArgs := services.TLSArgs{} + // tlsArgs := services.TLSArgs{} udpArgs := services.UDPArgs{} //build srvice args @@ -31,8 +31,8 @@ func initConfig() (err error) { args.Local = app.Flag("local", "local ip:port to listen").Short('p').Default(":33080").String() certTLS := app.Flag("cert", "cert file for tls").Short('C').Default("proxy.crt").String() keyTLS := app.Flag("key", "key file for tls").Short('K').Default("proxy.key").String() - args.PoolSize = app.Flag("pool-size", "conn pool size , which connect to parent proxy, zero: means turn off pool").Default("50").Int() - args.CheckParentInterval = app.Flag("check-parent-interval", "check if proxy is okay every interval seconds,zero: means no check").Default("3").Int() + args.PoolSize = app.Flag("pool-size", "conn pool size , which connect to parent proxy, zero: means turn off pool").Short('L').Default("50").Int() + args.CheckParentInterval = app.Flag("check-parent-interval", "check if proxy is okay every interval seconds,zero: means no check").Short('I').Default("3").Int() //########http######### http := app.Command("http", "proxy on http mode") @@ -49,13 +49,13 @@ func initConfig() (err error) { //########tcp######### tcp := app.Command("tcp", "proxy on tcp mode") - tcpArgs.Timeout = tcp.Flag("timeout", "tcp timeout milliseconds when connect to real server or parent proxy").Default("2000").Int() + tcpArgs.Timeout = tcp.Flag("timeout", "tcp timeout milliseconds when connect to real server or parent proxy").Short('t').Default("2000").Int() tcpArgs.ParentType = tcp.Flag("parent-type", "parent protocol type ").Short('T').Enum("tls", "tcp", "udp") - - //########tls######### - tls := app.Command("tls", "proxy on tls mode") - tlsArgs.Timeout = tls.Flag("timeout", "tcp timeout milliseconds when connect to real server or parent proxy").Default("2000").Int() - tlsArgs.ParentType = tls.Flag("parent-type", "parent protocol type ").Short('T').Enum("tls", "tcp", "udp") + tcpArgs.IsTLS = tcp.Flag("tls", "proxy on tls mode").Default("false").Bool() + //########udp######### + udp := app.Command("udp", "proxy on udp mode") + udpArgs.Timeout = udp.Flag("timeout", "tcp timeout milliseconds when connect to parent proxy").Short('t').Default("2000").Int() + udpArgs.ParentType = udp.Flag("parent-type", "parent protocol type ").Short('T').Enum("tls", "tcp", "udp") kingpin.MustParse(app.Parse(os.Args[1:])) @@ -64,7 +64,7 @@ func initConfig() (err error) { } httpArgs.Args = args tcpArgs.Args = args - tlsArgs.Args = args + // tlsArgs.Args = args udpArgs.Args = args //keygen @@ -78,7 +78,6 @@ func initConfig() (err error) { serviceName := kingpin.MustParse(app.Parse(os.Args[1:])) services.Regist("http", services.NewHTTP(), httpArgs) services.Regist("tcp", services.NewTCP(), tcpArgs) - services.Regist("tls", services.NewTLS(), tlsArgs) services.Regist("udp", services.NewUDP(), udpArgs) service, err = services.Run(serviceName) if err != nil { diff --git a/services/args.go b/services/args.go index d33b119..17ba89d 100644 --- a/services/args.go +++ b/services/args.go @@ -22,12 +22,9 @@ type TCPArgs struct { Args Timeout *int ParentType *string + IsTLS *bool } -type TLSArgs struct { - Args - Timeout *int - ParentType *string -} + type HTTPArgs struct { Args Always *bool @@ -43,4 +40,13 @@ type HTTPArgs struct { } type UDPArgs struct { Args + ParentType *string + Timeout *int +} + +func (a *TCPArgs) Protocol() string { + if *a.IsTLS { + return "tls" + } + return "tcp" } diff --git a/services/http.go b/services/http.go index 332b9b3..7b5c472 100644 --- a/services/http.go +++ b/services/http.go @@ -2,6 +2,7 @@ package services import ( "fmt" + "io" "log" "net" "proxy/utils" @@ -61,39 +62,39 @@ func (s *HTTP) Clean() { s.StopService() } func (s *HTTP) callback(inConn net.Conn) { - go func() { - defer func() { - if err := recover(); err != nil { - log.Printf("http(s) conn handler crashed with err : %s \nstack: %s", err, string(debug.Stack())) - } - }() - req, err := utils.NewHTTPRequest(&inConn, 4096, s.IsBasicAuth(), &s.basicAuth) - if err != nil { - log.Printf("decoder error , form %s, ERR:%s", err, inConn.RemoteAddr()) - utils.CloseConn(&inConn) - return - } - address := req.Host - useProxy := true - if *s.cfg.Parent == "" { - useProxy = false - } else if *s.cfg.Always { - useProxy = true - } else { - useProxy, _, _ = s.checker.IsBlocked(req.Host) - } - log.Printf("use proxy : %v, %s", useProxy, address) - //os.Exit(0) - err = s.OutToTCP(useProxy, address, &inConn, &req) - if err != nil { - if *s.cfg.Parent == "" { - log.Printf("connect to %s fail, ERR:%s", address, err) - } else { - log.Printf("connect to %s parent %s fail", *s.cfg.ParentType, *s.cfg.Parent) - } - utils.CloseConn(&inConn) + defer func() { + if err := recover(); err != nil { + log.Printf("http(s) conn handler crashed with err : %s \nstack: %s", err, string(debug.Stack())) } }() + req, err := utils.NewHTTPRequest(&inConn, 4096, s.IsBasicAuth(), &s.basicAuth) + if err != nil { + if err != io.EOF { + log.Printf("decoder error , form %s, ERR:%s", err, inConn.RemoteAddr()) + } + utils.CloseConn(&inConn) + return + } + address := req.Host + useProxy := true + if *s.cfg.Parent == "" { + useProxy = false + } else if *s.cfg.Always { + useProxy = true + } else { + useProxy, _, _ = s.checker.IsBlocked(req.Host) + } + log.Printf("use proxy : %v, %s", useProxy, address) + //os.Exit(0) + err = s.OutToTCP(useProxy, address, &inConn, &req) + if err != nil { + if *s.cfg.Parent == "" { + log.Printf("connect to %s fail, ERR:%s", address, err) + } else { + log.Printf("connect to %s parent %s fail", *s.cfg.ParentType, *s.cfg.Parent) + } + utils.CloseConn(&inConn) + } } func (s *HTTP) OutToTCP(useProxy bool, address string, inConn *net.Conn, req *utils.HTTPRequest) (err error) { inAddr := (*inConn).RemoteAddr().String() diff --git a/services/service.go b/services/service.go index 575763a..285a504 100644 --- a/services/service.go +++ b/services/service.go @@ -16,16 +16,16 @@ type ServiceItem struct { Name string } -var servicesMap = map[string]ServiceItem{} +var servicesMap = map[string]*ServiceItem{} func Regist(name string, s Service, args interface{}) { - servicesMap[name] = ServiceItem{ + servicesMap[name] = &ServiceItem{ S: s, Args: args, Name: name, } } -func Run(name string) (service ServiceItem, err error) { +func Run(name string) (service *ServiceItem, err error) { service, ok := servicesMap[name] if ok { go func() { diff --git a/services/tcp.go b/services/tcp.go index 0c27a0e..7a2b487 100644 --- a/services/tcp.go +++ b/services/tcp.go @@ -2,10 +2,12 @@ package services import ( "fmt" + "io" "log" "net" "proxy/utils" "runtime/debug" + "time" "strconv" ) @@ -34,7 +36,7 @@ func (s *TCP) Start(args interface{}) (err error) { if *s.cfg.Parent != "" { log.Printf("use %s parent %s", *s.cfg.ParentType, *s.cfg.Parent) } else { - log.Fatalf("parent required for tcp", *s.cfg.Local) + log.Fatalf("parent required for %s %s", s.cfg.Protocol(), *s.cfg.Local) } s.InitService() @@ -42,41 +44,43 @@ func (s *TCP) Start(args interface{}) (err error) { host, port, _ := net.SplitHostPort(*s.cfg.Local) p, _ := strconv.Atoi(port) sc := utils.NewServerChannel(host, p) - err = sc.ListenTCP(func(inConn net.Conn) { - go func() { - defer func() { - if err := recover(); err != nil { - log.Printf("tcp conn handler crashed with err : %s \nstack: %s", err, string(debug.Stack())) - } - }() - var err error - switch *s.cfg.ParentType { - case TYPE_TCP: - fallthrough - case TYPE_TLS: - err = s.OutToTCP(&inConn) - case TYPE_UDP: - err = s.OutToUDP(&inConn) - default: - err = fmt.Errorf("unkown parent type %s", *s.cfg.ParentType) - } - if err != nil { - log.Printf("connect to %s parent %s fail, ERR:%s", *s.cfg.ParentType, *s.cfg.Parent, err) - utils.CloseConn(&inConn) - } - }() - }) + if !*s.cfg.IsTLS { + err = sc.ListenTCP(s.callback) + } else { + err = sc.ListenTls(s.cfg.CertBytes, s.cfg.KeyBytes, s.callback) + } if err != nil { return } - log.Printf("tcp proxy on %s", (*sc.Listener).Addr()) + log.Printf("%s proxy on %s", s.cfg.Protocol(), (*sc.Listener).Addr()) return } func (s *TCP) Clean() { s.StopService() } - +func (s *TCP) callback(inConn net.Conn) { + defer func() { + if err := recover(); err != nil { + log.Printf("%s conn handler crashed with err : %s \nstack: %s", s.cfg.Protocol(), err, string(debug.Stack())) + } + }() + var err error + switch *s.cfg.ParentType { + case TYPE_TCP: + fallthrough + case TYPE_TLS: + err = s.OutToTCP(&inConn) + case TYPE_UDP: + err = s.OutToUDP(&inConn) + default: + err = fmt.Errorf("unkown parent type %s", *s.cfg.ParentType) + } + if err != nil { + log.Printf("connect to %s parent %s fail, ERR:%s", *s.cfg.ParentType, *s.cfg.Parent, err) + utils.CloseConn(&inConn) + } +} func (s *TCP) OutToTCP(inConn *net.Conn) (err error) { var outConn net.Conn var _outConn interface{} @@ -102,7 +106,52 @@ func (s *TCP) OutToTCP(inConn *net.Conn) (err error) { return } func (s *TCP) OutToUDP(inConn *net.Conn) (err error) { + log.Printf("conn created , remote : %s ", (*inConn).RemoteAddr()) + for { + srcAddr, body, err := utils.ReadUDPPacket(inConn) + if err == io.EOF || err == io.ErrUnexpectedEOF { + //log.Printf("connection %s released", srcAddr) + utils.CloseConn(inConn) + break + } + //log.Debugf("udp packet revecived:%s,%v", srcAddr, body) + dstAddr, err := net.ResolveUDPAddr("udp", *s.cfg.Parent) + if err != nil { + log.Printf("can't resolve address: %s", err) + utils.CloseConn(inConn) + break + } + clientSrcAddr := &net.UDPAddr{IP: net.IPv4zero, Port: 0} + conn, err := net.DialUDP("udp", clientSrcAddr, dstAddr) + if err != nil { + log.Printf("connect to udp %s fail,ERR:%s", dstAddr.String(), err) + continue + } + conn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout))) + _, err = conn.Write(body) + if err != nil { + log.Printf("send udp packet to %s fail,ERR:%s", dstAddr.String(), err) + continue + } + //log.Debugf("send udp packet to %s success", dstAddr.String()) + buf := make([]byte, 512) + len, _, err := conn.ReadFromUDP(buf) + if err != nil { + log.Printf("read udp response from %s fail ,ERR:%s", dstAddr.String(), err) + continue + } + respBody := buf[0:len] + //log.Debugf("revecived udp packet from %s , %v", dstAddr.String(), respBody) + _, err = (*inConn).Write(utils.UDPPacket(srcAddr, respBody)) + if err != nil { + log.Printf("send udp response fail ,ERR:%s", err) + utils.CloseConn(inConn) + break + } + //log.Printf("send udp response success ,from:%s", dstAddr.String()) + } return + } func (s *TCP) InitOutConnPool() { if *s.cfg.ParentType == TYPE_TLS || *s.cfg.ParentType == TYPE_TCP { diff --git a/services/tls.go b/services/tls.go deleted file mode 100644 index 2c09dc0..0000000 --- a/services/tls.go +++ /dev/null @@ -1,117 +0,0 @@ -package services - -import ( - "fmt" - "log" - "net" - "proxy/utils" - "runtime/debug" - "strconv" -) - -type TLS struct { - outPool utils.OutPool - cfg TLSArgs -} - -func NewTLS() Service { - return &TLS{} -} -func (s *TLS) InitService() { - s.InitOutConnPool() -} -func (s *TLS) StopService() { - if s.outPool.Pool != nil { - s.outPool.Pool.ReleaseAll() - } -} -func (s *TLS) Start(args interface{}) (err error) { - s.cfg = args.(TLSArgs) - if *s.cfg.Parent != "" { - log.Printf("use %s parent %s", *s.cfg.ParentType, *s.cfg.Parent) - } else { - log.Fatalf("parent required for tls") - } - - s.InitService() - - host, port, _ := net.SplitHostPort(*s.cfg.Local) - p, _ := strconv.Atoi(port) - sc := utils.NewServerChannel(host, p) - err = sc.ListenTls(s.cfg.CertBytes, s.cfg.KeyBytes, func(inConn net.Conn) { - go func() { - defer func() { - if err := recover(); err != nil { - log.Printf("tls conn handler crashed with err : %s \nstack: %s", err, string(debug.Stack())) - } - }() - var err error - switch *s.cfg.ParentType { - case TYPE_TCP: - fallthrough - case TYPE_TLS: - err = s.OutToTCP(&inConn) - case TYPE_UDP: - err = s.OutToUDP(&inConn) - default: - err = fmt.Errorf("unkown parent type %s", *s.cfg.ParentType) - } - if err != nil { - log.Printf("connect to %s parent %s fail, ERR:%s", *s.cfg.ParentType, *s.cfg.Parent, err) - utils.CloseConn(&inConn) - } - }() - }) - if err != nil { - return - } - log.Printf("tls proxy on %s", (*sc.Listener).Addr()) - return -} - -func (s *TLS) Clean() { - s.StopService() -} - -func (s *TLS) OutToTCP(inConn *net.Conn) (err error) { - var outConn net.Conn - var _outConn interface{} - _outConn, err = s.outPool.Pool.Get() - if err == nil { - outConn = _outConn.(net.Conn) - } - if err != nil { - log.Printf("connect to %s , err:%s", *s.cfg.Parent, err) - utils.CloseConn(inConn) - return - } - inAddr := (*inConn).RemoteAddr().String() - inLocalAddr := (*inConn).LocalAddr().String() - outAddr := outConn.RemoteAddr().String() - outLocalAddr := outConn.LocalAddr().String() - utils.IoBind((*inConn), outConn, func(err error) { - log.Printf("conn %s - %s - %s -%s released", inAddr, inLocalAddr, outLocalAddr, outAddr) - utils.CloseConn(inConn) - utils.CloseConn(&outConn) - }, func(n int, d bool) {}, 0) - log.Printf("conn %s - %s - %s -%s connected", inAddr, inLocalAddr, outLocalAddr, outAddr) - return -} -func (s *TLS) OutToUDP(inConn *net.Conn) (err error) { - return -} -func (s *TLS) InitOutConnPool() { - if *s.cfg.ParentType == TYPE_TLS || *s.cfg.ParentType == TYPE_TCP { - //dur int, isTLS bool, certBytes, keyBytes []byte, - //parent string, timeout int, InitialCap int, MaxCap int - s.outPool = utils.NewOutPool( - *s.cfg.CheckParentInterval, - *s.cfg.ParentType == TYPE_TLS, - s.cfg.CertBytes, s.cfg.KeyBytes, - *s.cfg.Parent, - *s.cfg.Timeout, - *s.cfg.PoolSize, - *s.cfg.PoolSize*2, - ) - } -} diff --git a/services/udp.go b/services/udp.go index 9f9aaa0..2850f86 100644 --- a/services/udp.go +++ b/services/udp.go @@ -1,19 +1,207 @@ package services import ( + "bufio" + "fmt" + "hash/crc32" + "io" "log" + "net" + "proxy/utils" + "runtime/debug" + "strconv" + "strings" + "time" ) type UDP struct { + p utils.ConcurrentMap + outPool utils.OutPool + cfg UDPArgs + sc *utils.ServerChannel } func NewUDP() Service { - return &UDP{} + return &UDP{ + outPool: utils.OutPool{}, + p: utils.NewConcurrentMap(), + } +} +func (s *UDP) InitService() { + if *s.cfg.ParentType != TYPE_UDP { + s.InitOutConnPool() + } +} +func (s *UDP) StopService() { + if s.outPool.Pool != nil { + s.outPool.Pool.ReleaseAll() + } } func (s *UDP) Start(args interface{}) (err error) { - log.Printf("called") + s.cfg = args.(UDPArgs) + if *s.cfg.Parent != "" { + log.Printf("use %s parent %s", *s.cfg.ParentType, *s.cfg.Parent) + } else { + log.Fatalf("parent required for udp %s", *s.cfg.Local) + } + + s.InitService() + + host, port, _ := net.SplitHostPort(*s.cfg.Local) + p, _ := strconv.Atoi(port) + sc := utils.NewServerChannel(host, p) + s.sc = &sc + err = sc.ListenUDP(s.callback) + if err != nil { + return + } + log.Printf("udp proxy on %s", (*sc.UDPListener).LocalAddr()) return } -func (s *UDP) Clean() { +func (s *UDP) Clean() { + s.StopService() +} +func (s *UDP) callback(packet []byte, localAddr, srcAddr *net.UDPAddr) { + defer func() { + if err := recover(); err != nil { + log.Printf("udp conn handler crashed with err : %s \nstack: %s", err, string(debug.Stack())) + } + }() + var err error + switch *s.cfg.ParentType { + case TYPE_TCP: + fallthrough + case TYPE_TLS: + err = s.OutToTCP(packet, localAddr, srcAddr) + case TYPE_UDP: + err = s.OutToUDP(packet, localAddr, srcAddr) + default: + err = fmt.Errorf("unkown parent type %s", *s.cfg.ParentType) + } + if err != nil { + log.Printf("connect to %s parent %s fail, ERR:%s", *s.cfg.ParentType, *s.cfg.Parent, err) + } +} +func (s *UDP) GetConn(connKey string) (conn net.Conn, isNew bool, err error) { + isNew = !s.p.Has(connKey) + var _conn interface{} + if isNew { + _conn, err = s.outPool.Pool.Get() + if err != nil { + return nil, false, err + } + s.p.Set(connKey, _conn) + } else { + _conn, _ = s.p.Get(connKey) + } + conn = _conn.(net.Conn) + return +} +func (s *UDP) OutToTCP(packet []byte, localAddr, srcAddr *net.UDPAddr) (err error) { + numLocal := crc32.ChecksumIEEE([]byte(localAddr.String())) + numSrc := crc32.ChecksumIEEE([]byte(srcAddr.String())) + connKey := uint64((numLocal/10)*10 + numSrc%10) + conn, isNew, err := s.GetConn(fmt.Sprintf("%d", connKey)) + if err != nil { + log.Printf("upd get conn to %s parent %s fail, ERR:%s", *s.cfg.ParentType, *s.cfg.Parent, err) + return + } + if isNew { + go func() { + defer func() { + if err := recover(); err != nil { + log.Printf("udp conn handler out to tcp crashed with err : %s \nstack: %s", err, string(debug.Stack())) + } + }() + log.Printf("conn %d created , local: %s", connKey, srcAddr.String()) + for { + srcAddrFromConn, body, err := utils.ReadUDPPacket(&conn) + if err == io.EOF || err == io.ErrUnexpectedEOF { + //log.Printf("connection %d released", connKey) + s.p.Remove(fmt.Sprintf("%d", connKey)) + break + } + if err != nil { + log.Printf("parse revecived udp packet fail, err: %s", err) + continue + } + //log.Printf("udp packet revecived over parent , local:%s", srcAddrFromConn) + _srcAddr := strings.Split(srcAddrFromConn, ":") + if len(_srcAddr) != 2 { + log.Printf("parse revecived udp packet fail, addr error : %s", srcAddrFromConn) + continue + } + port, _ := strconv.Atoi(_srcAddr[1]) + dstAddr := &net.UDPAddr{IP: net.ParseIP(_srcAddr[0]), Port: port} + _, err = s.sc.UDPListener.WriteToUDP(body, dstAddr) + if err != nil { + log.Printf("udp response to local %s fail,ERR:%s", srcAddr, err) + continue + } + //log.Printf("udp response to local %s success", srcAddr) + } + }() + } + //log.Printf("select conn %d , local: %s", connKey, srcAddr.String()) + writer := bufio.NewWriter(conn) + //fmt.Println(conn, writer) + writer.Write(utils.UDPPacket(srcAddr.String(), packet)) + err = writer.Flush() + if err != nil { + log.Printf("write udp packet to %s fail ,flush err:%s", *s.cfg.Parent, err) + return + } + //log.Printf("write packet %v", packet) + return +} +func (s *UDP) OutToUDP(packet []byte, localAddr, srcAddr *net.UDPAddr) (err error) { + //log.Printf("udp packet revecived:%s,%v", srcAddr, packet) + dstAddr, err := net.ResolveUDPAddr("udp", *s.cfg.Parent) + if err != nil { + log.Printf("resolve udp addr %s fail fail,ERR:%s", dstAddr.String(), err) + return + } + clientSrcAddr := &net.UDPAddr{IP: net.IPv4zero, Port: 0} + conn, err := net.DialUDP("udp", clientSrcAddr, dstAddr) + if err != nil { + log.Printf("connect to udp %s fail,ERR:%s", dstAddr.String(), err) + return + } + conn.SetDeadline(time.Now().Add(time.Millisecond * time.Duration(*s.cfg.Timeout))) + _, err = conn.Write(packet) + if err != nil { + log.Printf("send udp packet to %s fail,ERR:%s", dstAddr.String(), err) + return + } + //log.Printf("send udp packet to %s success", dstAddr.String()) + buf := make([]byte, 512) + len, _, err := conn.ReadFromUDP(buf) + if err != nil { + log.Printf("read udp response from %s fail ,ERR:%s", dstAddr.String(), err) + return + } + //log.Printf("revecived udp packet from %s , %v", dstAddr.String(), respBody) + _, err = s.sc.UDPListener.WriteToUDP(buf[0:len], srcAddr) + if err != nil { + log.Printf("send udp response to cluster fail ,ERR:%s", err) + return + } + //log.Printf("send udp response to cluster success ,from:%s", dstAddr.String()) + return +} +func (s *UDP) InitOutConnPool() { + if *s.cfg.ParentType == TYPE_TLS || *s.cfg.ParentType == TYPE_TCP { + //dur int, isTLS bool, certBytes, keyBytes []byte, + //parent string, timeout int, InitialCap int, MaxCap int + s.outPool = utils.NewOutPool( + *s.cfg.CheckParentInterval, + *s.cfg.ParentType == TYPE_TLS, + s.cfg.CertBytes, s.cfg.KeyBytes, + *s.cfg.Parent, + *s.cfg.Timeout, + *s.cfg.PoolSize, + *s.cfg.PoolSize*2, + ) + } } diff --git a/utils/functions.go b/utils/functions.go index 9eb8ed1..31f8277 100755 --- a/utils/functions.go +++ b/utils/functions.go @@ -1,8 +1,11 @@ package utils import ( + "bufio" + "bytes" "crypto/tls" "crypto/x509" + "encoding/binary" "errors" "fmt" "io" @@ -254,6 +257,49 @@ func GetAllInterfaceAddr() ([]net.IP, error) { //only need first return addresses, nil } +func UDPPacket(srcAddr string, packet []byte) []byte { + addrBytes := []byte(srcAddr) + addrLength := uint16(len(addrBytes)) + bodyLength := uint16(len(packet)) + pkg := new(bytes.Buffer) + binary.Write(pkg, binary.LittleEndian, addrLength) + binary.Write(pkg, binary.LittleEndian, addrBytes) + binary.Write(pkg, binary.LittleEndian, bodyLength) + binary.Write(pkg, binary.LittleEndian, packet) + return pkg.Bytes() +} +func ReadUDPPacket(conn *net.Conn) (srcAddr string, packet []byte, err error) { + reader := bufio.NewReader(*conn) + var addrLength uint16 + var bodyLength uint16 + err = binary.Read(reader, binary.LittleEndian, &addrLength) + if err != nil { + return + } + _srcAddr := make([]byte, addrLength) + n, err := reader.Read(_srcAddr) + if err != nil { + return + } + if n != int(addrLength) { + return + } + srcAddr = string(_srcAddr) + + err = binary.Read(reader, binary.LittleEndian, &bodyLength) + if err != nil { + return + } + packet = make([]byte, bodyLength) + n, err = reader.Read(packet) + if err != nil { + return + } + if n != int(bodyLength) { + return + } + return +} // type sockaddr struct { // family uint16