Przeglądaj źródła

icmp: simplify Message.Marshal

Change-Id: I1c78c929dd824534ccc2123a63973af89b8a0081
Reviewed-on: https://go-review.googlesource.com/c/155857
Reviewed-by: Ian Lance Taylor <iant@golang.org>
Mikio Hara 7 lat temu
rodzic
commit
4602d50750
1 zmienionych plików z 9 dodań i 8 usunięć
  1. 9 8
      icmp/message.go

+ 9 - 8
icmp/message.go

@@ -75,27 +75,28 @@ type Message struct {
 // compute the checksum field during the message transmission.
 // compute the checksum field during the message transmission.
 // When psh is not nil, it must be the pseudo header for IPv6.
 // When psh is not nil, it must be the pseudo header for IPv6.
 func (m *Message) Marshal(psh []byte) ([]byte, error) {
 func (m *Message) Marshal(psh []byte) ([]byte, error) {
-	var mtype int
+	var mtype byte
 	switch typ := m.Type.(type) {
 	switch typ := m.Type.(type) {
 	case ipv4.ICMPType:
 	case ipv4.ICMPType:
-		mtype = int(typ)
+		mtype = byte(typ)
 	case ipv6.ICMPType:
 	case ipv6.ICMPType:
-		mtype = int(typ)
+		mtype = byte(typ)
 	default:
 	default:
 		return nil, errInvalidProtocol
 		return nil, errInvalidProtocol
 	}
 	}
-	b := []byte{byte(mtype), byte(m.Code), 0, 0}
-	if m.Type.Protocol() == iana.ProtocolIPv6ICMP && psh != nil {
+	b := []byte{mtype, byte(m.Code), 0, 0}
+	proto := m.Type.Protocol()
+	if proto == iana.ProtocolIPv6ICMP && psh != nil {
 		b = append(psh, b...)
 		b = append(psh, b...)
 	}
 	}
-	if m.Body != nil && m.Body.Len(m.Type.Protocol()) != 0 {
-		mb, err := m.Body.Marshal(m.Type.Protocol())
+	if m.Body != nil && m.Body.Len(proto) != 0 {
+		mb, err := m.Body.Marshal(proto)
 		if err != nil {
 		if err != nil {
 			return nil, err
 			return nil, err
 		}
 		}
 		b = append(b, mb...)
 		b = append(b, mb...)
 	}
 	}
-	if m.Type.Protocol() == iana.ProtocolIPv6ICMP {
+	if proto == iana.ProtocolIPv6ICMP {
 		if psh == nil { // cannot calculate checksum here
 		if psh == nil { // cannot calculate checksum here
 			return b, nil
 			return b, nil
 		}
 		}