p2p: fix disconnection encoding on the writing side + fix decoding during protocol handshake

This commit is contained in:
lorenzo 2024-12-10 09:52:30 +01:00
parent 78b3b9206d
commit 5a2b9bde72

View file

@ -114,12 +114,14 @@ func (t *rlpxTransport) close(err error) {
// We only bother doing this if the underlying connection supports // We only bother doing this if the underlying connection supports
// setting a timeout tough. // setting a timeout tough.
if t.conn != nil { if t.conn != nil {
if r, ok := err.(DiscReason); ok && r != DiscNetworkError { if reason, ok := err.(DiscReason); ok && reason != DiscNetworkError {
// We do not use the WriteMsg func since we want a custom deadline
deadline := time.Now().Add(discWriteTimeout) deadline := time.Now().Add(discWriteTimeout)
if err := t.conn.SetWriteDeadline(deadline); err == nil { if err := t.conn.SetWriteDeadline(deadline); err == nil {
// Connection supports write deadline. // Connection supports write deadline.
t.wbuf.Reset() t.wbuf.Reset()
rlp.Encode(&t.wbuf, []DiscReason{r}) size, reader, _ := rlp.EncodeToReader(reason)
io.CopyN(&t.wbuf, reader, int64(size))
t.conn.Write(discMsg, t.wbuf.Bytes()) t.conn.Write(discMsg, t.wbuf.Bytes())
} }
} }
@ -165,9 +167,9 @@ func readProtocolHandshake(rw MsgReader) (*protoHandshake, error) {
// spec and we send it ourself if the post-handshake checks fail. // spec and we send it ourself if the post-handshake checks fail.
// We can't return the reason directly, though, because it is echoed // We can't return the reason directly, though, because it is echoed
// back otherwise. Wrap it in a string instead. // back otherwise. Wrap it in a string instead.
var reason [1]DiscReason var m struct{ R DiscReason }
rlp.Decode(msg.Payload, &reason) rlp.Decode(msg.Payload, &m)
return nil, reason[0] return nil, m.R
} }
if msg.Code != handshakeMsg { if msg.Code != handshakeMsg {
return nil, fmt.Errorf("expected handshake, got %x", msg.Code) return nil, fmt.Errorf("expected handshake, got %x", msg.Code)