feat: support diag TCP with only socket returns

For something like statistics that just requires socket info only
This commit is contained in:
Neo Zhuo 2021-05-27 18:27:05 +08:00 committed by Alessandro Boch
parent f9d8f911db
commit 14e832ae1e

View File

@ -172,12 +172,56 @@ func SocketGet(local, remote net.Addr) (*Socket, error) {
return sock, nil
}
// SocketDiagTCPInfo requests INET_DIAG_INFO for TCP protocol for specified family type.
// SocketDiagTCPInfo requests INET_DIAG_INFO for TCP protocol for specified family type and return with extension TCP info.
func SocketDiagTCPInfo(family uint8) ([]*InetDiagTCPInfoResp, error) {
s, err := nl.Subscribe(unix.NETLINK_INET_DIAG)
var result []*InetDiagTCPInfoResp
err := socketDiagTCPExecutor(family, func(m syscall.NetlinkMessage) error {
sockInfo := &Socket{}
if err := sockInfo.deserialize(m.Data); err != nil {
return err
}
attrs, err := nl.ParseRouteAttr(m.Data[sizeofSocket:])
if err != nil {
return err
}
res, err := attrsToInetDiagTCPInfoResp(attrs, sockInfo)
if err != nil {
return err
}
result = append(result, res)
return nil
})
if err != nil {
return nil, err
}
return result, nil
}
// SocketDiagTCP requests INET_DIAG_INFO for TCP protocol for specified family type and return related socket.
func SocketDiagTCP(family uint8) ([]*Socket, error) {
var result []*Socket
err := socketDiagTCPExecutor(family, func(m syscall.NetlinkMessage) error {
sockInfo := &Socket{}
if err := sockInfo.deserialize(m.Data); err != nil {
return err
}
result = append(result, sockInfo)
return nil
})
if err != nil {
return nil, err
}
return result, nil
}
// socketDiagTCPExecutor requests INET_DIAG_INFO for TCP protocol for specified family type.
func socketDiagTCPExecutor(family uint8, receiver func(syscall.NetlinkMessage) error) error {
s, err := nl.Subscribe(unix.NETLINK_INET_DIAG)
if err != nil {
return err
}
defer s.Close()
req := nl.NewNetlinkRequest(nl.SOCK_DIAG_BY_FAMILY, unix.NLM_F_DUMP)
@ -189,18 +233,17 @@ func SocketDiagTCPInfo(family uint8) ([]*InetDiagTCPInfoResp, error) {
})
s.Send(req)
var result []*InetDiagTCPInfoResp
loop:
for {
msgs, from, err := s.Receive()
if err != nil {
return nil, err
return err
}
if from.Pid != nl.PidKernel {
return nil, fmt.Errorf("Wrong sender portid %d, expected %d", from.Pid, nl.PidKernel)
return fmt.Errorf("Wrong sender portid %d, expected %d", from.Pid, nl.PidKernel)
}
if len(msgs) == 0 {
return nil, errors.New("no message nor error from netlink")
return errors.New("no message nor error from netlink")
}
for _, m := range msgs {
@ -209,26 +252,14 @@ loop:
break loop
case unix.NLMSG_ERROR:
error := int32(native.Uint32(m.Data[0:4]))
return nil, syscall.Errno(-error)
return syscall.Errno(-error)
}
sockInfo := &Socket{}
if err := sockInfo.deserialize(m.Data); err != nil {
return nil, err
if err := receiver(m); err != nil {
return err
}
attrs, err := nl.ParseRouteAttr(m.Data[sizeofSocket:])
if err != nil {
return nil, err
}
res, err := attrsToInetDiagTCPInfoResp(attrs, sockInfo)
if err != nil {
return nil, err
}
result = append(result, res)
}
}
return result, nil
return nil
}
func attrsToInetDiagTCPInfoResp(attrs []syscall.NetlinkRouteAttr, sockInfo *Socket) (*InetDiagTCPInfoResp, error) {