go简易tcp/udp连接测试工具

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

package main
 
import (
    "fmt"
    "io"
    "log"
    "net"
    "os"
)
 
func showBytes(bytes []byte) string {
    var str []byte
    for _, b := range bytes {
        if b >= 33 && b <= 126 && b != 92 {
            str = append(str, b)
        } else {
            str = append(str, fmt.Sprintf("\\%02X", b)...)
        }
    }
    return string(str)
}
 
func writeAllTCP(con *net.TCPConn, message []byte) error {
    defer con.CloseWrite()
    for len(message) > 0 {
        l, err := con.Write(message)
        if err != nil {
            return err
        }
        message = message[l:]
    }
    return nil
}
 
func handleTCP(client *net.TCPConn, message []byte) {
    defer client.Close()
    addrStr := client.RemoteAddr().String()
    log.Printf("Client %s connected\n", addrStr)
    var data [1024]byte
    n, err := io.ReadFull(client, data[:])
    if err != nil && err != io.ErrUnexpectedEOF {
        log.Printf("Failed to read: %v\n", err)
        return
    }
    log.Printf("From %s: %s\n", addrStr, showBytes(data[:n]))
    if err := writeAllTCP(client, message); err != nil {
        log.Printf("Failed to write: %v\n", err)
        return
    }
}
 
func main() {
    if len(os.Args) != 5 || os.Args[1] != "tcp" && os.Args[1] != "udp" || os.Args[2] != "server" && os.Args[2] != "client" {
        fmt.Fprintf(os.Stderr, "Usage: %s tcp|udp server|client <server address> <message>\n", os.Args[0])
        os.Exit(1)
    }
    if os.Args[1] == "tcp" {
        addr, err := net.ResolveTCPAddr("tcp", os.Args[3])
        if err != nil {
            panic(err)
        }
        if os.Args[2] == "server" {
            con, err := net.ListenTCP("tcp", addr)
            if err != nil {
                panic(err)
            }
            defer con.Close()
            for {
                client, err := con.AcceptTCP()
                if err != nil {
                    log.Printf("Failed to accept: %v\n", err)
                    continue
                }
                go handleTCP(client, []byte(os.Args[4]))
            }
        } else {
            con, err := net.DialTCP("tcp", nil, addr)
            if err != nil {
                panic(err)
            }
            defer con.Close()
            if err := writeAllTCP(con, []byte(os.Args[4])); err != nil {
                panic(err)
            }
            var data [1024]byte
            n, err := io.ReadFull(con, data[:])
            if err != nil && err != io.ErrUnexpectedEOF {
                panic(err)
            }
            fmt.Println(showBytes(data[:n]))
        }
    } else {
        addr, err := net.ResolveUDPAddr("udp", os.Args[3])
        if err != nil {
            panic(err)
        }
        if os.Args[2] == "server" {
            con, err := net.ListenUDP("udp", addr)
            if err != nil {
                panic(err)
            }
            defer con.Close()
            for {
                var data [1024]byte
                n, addr, err := con.ReadFromUDP(data[:])
                if err != nil {
                    log.Printf("Failed to read: %v\n", err)
                    continue
                }
                log.Printf("From %v: %s\n", addr, showBytes(data[:n]))
                _, err = con.WriteTo([]byte(os.Args[4]), addr)
                if err != nil {
                    log.Printf("Failed to write: %v\n", err)
                    continue
                }
            }
        } else {
            con, err := net.DialUDP("udp", nil, addr)
            if err != nil {
                panic(err)
            }
            defer con.Close()
            _, err = con.Write([]byte(os.Args[4]))
            if err != nil {
                panic(err)
            }
            var data [1024]byte
            n, _, err := con.ReadFromUDP(data[:])
            if err != nil {
                panic(err)
            }
            fmt.Println(showBytes(data[:n]))
        }
    }
}