Respect SO_SNDBUF for network datagram endpoints

Previously, SO_SNDBUF was effectively a no-op. The stack should make
sure that only SO_SNDBUF bytes are ever in-flight for any given
socket/endpoint.

Fuchsia Bug: https://fxbug.dev/99070

PiperOrigin-RevId: 446792223

pkg/tcpip/stack/packet_buffer.go

@@ -53,6 +53,10 @@ type PacketBufferOptions struct {
 	// IsForwardedPacket identifies that the PacketBuffer being created is for a
 	// forwarded packet.
 	IsForwardedPacket bool
+
+	// OnRelease is a function to be run when the packet buffer is no longer
+	// referenced (released back to the pool).
+	OnRelease func()
 }
 
 // A PacketBuffer contains all the data of a network packet.
@@ -163,6 +167,10 @@ type PacketBuffer struct {
 	NetworkPacketInfo NetworkPacketInfo
 
 	tuple *tuple
+
+	// onRelease is a function to be run when the packet buffer is no longer
+	// referenced (released back to the pool).
+	onRelease func() `state:"nosave"`
 }
 
 // NewPacketBuffer creates a new PacketBuffer with opts.
@@ -177,6 +185,7 @@ func NewPacketBuffer(opts PacketBufferOptions) *PacketBuffer {
 		pk.buf.AppendOwned(v)
 	}
 	pk.NetworkPacketInfo.IsForwardedPacket = opts.IsForwardedPacket
+	pk.onRelease = opts.OnRelease
 	pk.InitRefs()
 	return pk
 }
@@ -186,6 +195,10 @@ func NewPacketBuffer(opts PacketBufferOptions) *PacketBuffer {
 // pool.
 func (pk *PacketBuffer) DecRef() {
 	pk.packetBufferRefs.DecRef(func() {
+		if pk.onRelease != nil {
+			pk.onRelease()
+		}
+
 		pkPool.Put(pk)
 	})
 }

pkg/tcpip/transport/BUILD

@@ -19,6 +19,7 @@ go_test(
     deps = [
         ":transport",
         "//pkg/tcpip",
+        "//pkg/tcpip/buffer",
         "//pkg/tcpip/header",
         "//pkg/tcpip/link/loopback",
         "//pkg/tcpip/network/ipv4",

pkg/tcpip/transport/datagram_test.go

@@ -16,10 +16,13 @@
 package transport_test
 
 import (
+	"bytes"
 	"fmt"
+	"math"
 	"testing"
 
 	"gvisor.dev/gvisor/pkg/tcpip"
+	"gvisor.dev/gvisor/pkg/tcpip/buffer"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
 	"gvisor.dev/gvisor/pkg/tcpip/link/loopback"
 	"gvisor.dev/gvisor/pkg/tcpip/network/ipv4"
@@ -126,3 +129,222 @@ func TestStateUpdates(t *testing.T) {
 		})
 	}
 }
+
+type mockEndpoint struct {
+	disp stack.NetworkDispatcher
+	pkts stack.PacketBufferList
+}
+
+func (*mockEndpoint) MTU() uint32 {
+	return math.MaxUint32
+}
+func (*mockEndpoint) Capabilities() stack.LinkEndpointCapabilities {
+	return 0
+}
+func (*mockEndpoint) MaxHeaderLength() uint16 {
+	return 0
+}
+func (*mockEndpoint) LinkAddress() tcpip.LinkAddress {
+	var l tcpip.LinkAddress
+	return l
+}
+func (e *mockEndpoint) WritePackets(pkts stack.PacketBufferList) (int, tcpip.Error) {
+	pkts.IncRef()
+	len := pkts.Len()
+	for pkt := pkts.Front(); pkt != nil; pkt = pkt.Next() {
+		e.pkts.PushBack(pkt)
+	}
+
+	return len, nil
+}
+func (e *mockEndpoint) Attach(d stack.NetworkDispatcher)      { e.disp = d }
+func (e *mockEndpoint) IsAttached() bool                      { return e.disp != nil }
+func (*mockEndpoint) Wait()                                   {}
+func (*mockEndpoint) ARPHardwareType() header.ARPHardwareType { return header.ARPHardwareNone }
+func (*mockEndpoint) AddHeader(*stack.PacketBuffer)           {}
+func (e *mockEndpoint) releasePackets() {
+	e.pkts.DecRef()
+	e.pkts = stack.PacketBufferList{}
+}
+
+func (e *mockEndpoint) pktsSize() int {
+	s := 0
+	for pkt := e.pkts.Front(); pkt != nil; pkt = pkt.Next() {
+		s += pkt.Size() + pkt.AvailableHeaderBytes()
+	}
+	return s
+}
+
+func TestSndBuf(t *testing.T) {
+	const nicID = 1
+
+	buf := buffer.NewView(header.ICMPv4MinimumSize)
+	header.ICMPv4(buf).SetType(header.ICMPv4Echo)
+
+	for _, test := range []struct {
+		name           string
+		createEndpoint func(*stack.Stack, *waiter.Queue) (tcpip.Endpoint, error)
+	}{
+		{
+			name: "UDP",
+			createEndpoint: func(s *stack.Stack, wq *waiter.Queue) (tcpip.Endpoint, error) {
+				ep, err := s.NewEndpoint(udp.ProtocolNumber, ipv4.ProtocolNumber, wq)
+				if err != nil {
+					return nil, fmt.Errorf("s.NewEndpoint(%d, %d, _) failed: %s", udp.ProtocolNumber, ipv4.ProtocolNumber, err)
+				}
+				return ep, nil
+			},
+		},
+		{
+			name: "ICMP",
+			createEndpoint: func(s *stack.Stack, wq *waiter.Queue) (tcpip.Endpoint, error) {
+				ep, err := s.NewEndpoint(icmp.ProtocolNumber4, ipv4.ProtocolNumber, wq)
+				if err != nil {
+					return nil, fmt.Errorf("s.NewEndpoint(%d, %d, _) failed: %s", icmp.ProtocolNumber4, ipv4.ProtocolNumber, err)
+				}
+				return ep, nil
+			},
+		},
+		{
+			name: "RAW",
+			createEndpoint: func(s *stack.Stack, wq *waiter.Queue) (tcpip.Endpoint, error) {
+				ep, err := s.NewRawEndpoint(udp.ProtocolNumber, ipv4.ProtocolNumber, wq, true /* associated */)
+				if err != nil {
+					return nil, fmt.Errorf("s.NewRawEndpoint(%d, %d, _, true) failed: %s", udp.ProtocolNumber, ipv4.ProtocolNumber, err)
+				}
+				return ep, nil
+			},
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			s := stack.New(stack.Options{
+				NetworkProtocols:   []stack.NetworkProtocolFactory{ipv4.NewProtocol},
+				TransportProtocols: []stack.TransportProtocolFactory{udp.NewProtocol, icmp.NewProtocol4},
+				RawFactory:         &raw.EndpointFactory{},
+			})
+			var e mockEndpoint
+			defer e.releasePackets()
+			if err := s.CreateNIC(nicID, &e); err != nil {
+				t.Fatalf("s.CreateNIC(%d, _) failed: %s", nicID, err)
+			}
+			var wq waiter.Queue
+			ep, err := test.createEndpoint(s, &wq)
+			if err != nil {
+				t.Fatalf("test.createEndpoint(_) failed: %s", err)
+			}
+			defer ep.Close()
+
+			addr := tcpip.ProtocolAddress{
+				Protocol:          ipv4.ProtocolNumber,
+				AddressWithPrefix: testutil.MustParse4("1.2.3.4").WithPrefix(),
+			}
+			if err := s.AddProtocolAddress(nicID, addr, stack.AddressProperties{}); err != nil {
+				t.Fatalf("AddProtocolAddress(%d, %#v, {}): %s", nicID, addr, err)
+			}
+			s.SetRouteTable([]tcpip.Route{
+				{
+					Destination: header.IPv4EmptySubnet,
+					NIC:         nicID,
+				},
+			})
+
+			to := tcpip.FullAddress{NIC: nicID, Addr: testutil.MustParse4("1.0.0.1"), Port: 12345}
+			if err := ep.Connect(to); err != nil {
+				t.Fatalf("ep.Connect(%#v): %s", to, err)
+			}
+
+			checkWriteFail := func() {
+				t.Helper()
+
+				if got := ep.Readiness(waiter.WritableEvents); got != 0 {
+					t.Fatalf("got ep.Readiness(0x%x) = 0x%x, want = 0x0", waiter.WritableEvents, got)
+				}
+
+				var r bytes.Reader
+				r.Reset(buf[:])
+				wantErr := &tcpip.ErrWouldBlock{}
+				if n, err := ep.Write(&r, tcpip.WriteOptions{}); err != wantErr {
+					t.Fatalf("got Write(...) = (%d, %s), want = (_, %s)", n, err, wantErr)
+				}
+
+			}
+
+			checkWrites := func() {
+				t.Helper()
+
+				if got := ep.Readiness(waiter.WritableEvents); got != waiter.WritableEvents {
+					t.Fatalf("got ep.Readiness(0x%x) = 0x%x, want = 0x%x", waiter.WritableEvents, got, waiter.WritableEvents)
+				}
+
+				var r bytes.Reader
+				r.Reset(buf[:])
+				if n, err := ep.Write(&r, tcpip.WriteOptions{}); err != nil {
+					t.Fatalf("Write(...): %s", err)
+				} else if want := int64(len(buf)); n != want {
+					t.Fatalf("got Write(...) = %d, want = %d", n, want)
+				}
+
+				// The next write should fail since the packet we sent before
+				// is still held.
+				checkWriteFail()
+			}
+
+			we, ch := waiter.NewChannelEntry(waiter.WritableEvents)
+			wq.EventRegister(&we)
+			defer wq.EventUnregister(&we)
+
+			checkNoWritableEvent := func() {
+				t.Helper()
+
+				select {
+				case <-ch:
+					t.Fatal("unexpected writable event")
+				default:
+				}
+			}
+
+			checkWritableEvent := func() {
+				t.Helper()
+
+				select {
+				case <-ch:
+				default:
+					t.Fatal("expected writable event")
+				}
+			}
+
+			// As long as there is space in the send buffer, writes should succeed
+			// so a send buffer of 1 allows at max 1 in-flight packet.
+			ep.SocketOptions().SetSendBufferSize(1, true)
+			checkWritableEvent()
+			checkWrites()
+			checkNoWritableEvent()
+
+			// Increase the size of the send buffer but still be full.
+			inUseSize := int64(e.pktsSize())
+			checkNoWritableEvent()
+			ep.SocketOptions().SetSendBufferSize(inUseSize, true /* notify */)
+			checkNoWritableEvent()
+			checkWriteFail()
+
+			// Open up the send buffer by 1 byte.
+			checkNoWritableEvent()
+			ep.SocketOptions().SetSendBufferSize(inUseSize+1, true /* notify */)
+			checkWritableEvent()
+			checkWrites()
+
+			// We can resize the send buffer to a smaller size but it is still
+			// full so we can't write.
+			checkNoWritableEvent()
+			ep.SocketOptions().SetSendBufferSize(1, true /* notify */)
+			checkNoWritableEvent()
+			checkWriteFail()
+
+			// Releasing the packets should open up the send buffer for the next
+			// write.
+			e.releasePackets()
+			checkWritableEvent()
+			checkWrites()
+		})
+	}
+}

pkg/tcpip/transport/icmp/endpoint.go

@@ -91,7 +91,7 @@ func newEndpoint(s *stack.Stack, netProto tcpip.NetworkProtocolNumber, transProt
 	ep.ops.InitHandler(ep, ep.stack, tcpip.GetStackSendBufferLimits, tcpip.GetStackReceiveBufferLimits)
 	ep.ops.SetSendBufferSize(32*1024, false /* notify */)
 	ep.ops.SetReceiveBufferSize(32*1024, false /* notify */)
-	ep.net.Init(s, netProto, transProto, &ep.ops)
+	ep.net.Init(s, netProto, transProto, &ep.ops, waiterQueue)
 
 	// Override with stack defaults.
 	var ss tcpip.SendBufferSizeOption
@@ -105,6 +105,11 @@ func newEndpoint(s *stack.Stack, netProto tcpip.NetworkProtocolNumber, transProt
 	return ep, nil
 }
 
+// WakeupWriters implements tcpip.SocketOptionsHandler.
+func (e *endpoint) WakeupWriters() {
+	e.net.MaybeSignalWritable()
+}
+
 // UniqueID implements stack.TransportEndpoint.UniqueID.
 func (e *endpoint) UniqueID() uint64 {
 	return e.uniqueID
@@ -399,9 +404,10 @@ func send4(s *stack.Stack, ctx *network.WriteContext, ident uint16, data buffer.
 		return &tcpip.ErrInvalidEndpointState{}
 	}
 
-	pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
-		ReserveHeaderBytes: header.ICMPv4MinimumSize + int(maxHeaderLength),
-	})
+	pkt := ctx.TryNewPacketBuffer(header.ICMPv4MinimumSize+int(maxHeaderLength), buffer.VectorisedView{})
+	if pkt == nil {
+		return &tcpip.ErrWouldBlock{}
+	}
 	defer pkt.DecRef()
 
 	icmpv4 := header.ICMPv4(pkt.TransportHeader().Push(header.ICMPv4MinimumSize))
@@ -440,9 +446,10 @@ func send6(s *stack.Stack, ctx *network.WriteContext, ident uint16, data buffer.
 		return &tcpip.ErrInvalidEndpointState{}
 	}
 
-	pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
-		ReserveHeaderBytes: header.ICMPv6MinimumSize + int(maxHeaderLength),
-	})
+	pkt := ctx.TryNewPacketBuffer(header.ICMPv6MinimumSize+int(maxHeaderLength), buffer.VectorisedView{})
+	if pkt == nil {
+		return &tcpip.ErrWouldBlock{}
+	}
 	defer pkt.DecRef()
 
 	icmpv6 := header.ICMPv6(pkt.TransportHeader().Push(header.ICMPv6MinimumSize))
@@ -662,8 +669,11 @@ func (e *endpoint) GetRemoteAddress() (tcpip.FullAddress, tcpip.Error) {
 // Readiness returns the current readiness of the endpoint. For example, if
 // waiter.EventIn is set, the endpoint is immediately readable.
 func (e *endpoint) Readiness(mask waiter.EventMask) waiter.EventMask {
-	// The endpoint is always writable.
-	result := waiter.WritableEvents & mask
+	var result waiter.EventMask
+
+	if e.net.HasSendSpace() {
+		result |= waiter.WritableEvents & mask
+	}
 
 	// Determine if the endpoint is readable if requested.
 	if (mask & waiter.ReadableEvents) != 0 {

pkg/tcpip/transport/internal/network/BUILD

@@ -17,9 +17,11 @@ go_library(
         "//pkg/atomicbitops",
         "//pkg/sync",
         "//pkg/tcpip",
+        "//pkg/tcpip/buffer",
         "//pkg/tcpip/header",
         "//pkg/tcpip/stack",
         "//pkg/tcpip/transport",
+        "//pkg/waiter",
     ],
 )
 
@@ -44,6 +46,7 @@ go_test(
         "//pkg/tcpip/testutil",
         "//pkg/tcpip/transport",
         "//pkg/tcpip/transport/udp",
+        "//pkg/waiter",
         "@com_github_google_go_cmp//cmp:go_default_library",
     ],
 )

pkg/tcpip/transport/internal/network/endpoint.go

@@ -22,9 +22,11 @@ import (
 	"gvisor.dev/gvisor/pkg/atomicbitops"
 	"gvisor.dev/gvisor/pkg/sync"
 	"gvisor.dev/gvisor/pkg/tcpip"
+	"gvisor.dev/gvisor/pkg/tcpip/buffer"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 	"gvisor.dev/gvisor/pkg/tcpip/transport"
+	"gvisor.dev/gvisor/pkg/waiter"
 )
 
 // Endpoint is a datagram-based endpoint. It only supports sending datagrams to
@@ -33,10 +35,11 @@ import (
 // +stateify savable
 type Endpoint struct {
 	// The following fields must only be set once then never changed.
-	stack      *stack.Stack `state:"manual"`
-	ops        *tcpip.SocketOptions
-	netProto   tcpip.NetworkProtocolNumber
-	transProto tcpip.TransportProtocolNumber
+	stack       *stack.Stack `state:"manual"`
+	ops         *tcpip.SocketOptions
+	netProto    tcpip.NetworkProtocolNumber
+	transProto  tcpip.TransportProtocolNumber
+	waiterQueue *waiter.Queue
 
 	mu sync.RWMutex `state:"nosave"`
 	// +checklocks:mu
@@ -96,6 +99,14 @@ type Endpoint struct {
 	//
 	// Writes must be performed through setEndpointState.
 	state atomicbitops.Uint32
+
+	// Callers should not attempt to obtain sendBufferSizeInUseMu while holding
+	// another lock on Endpoint.
+	sendBufferSizeInUseMu sync.RWMutex `state:"nosave"`
+	// sendBufferSizeInUse keeps track of the bytes in use by in-flight packets.
+	//
+	// +checklocks:sendBufferSizeInUseMu
+	sendBufferSizeInUse int64 `state:"nosave"`
 }
 
 // +stateify savable
@@ -105,7 +116,7 @@ type multicastMembership struct {
 }
 
 // Init initializes the endpoint.
-func (e *Endpoint) Init(s *stack.Stack, netProto tcpip.NetworkProtocolNumber, transProto tcpip.TransportProtocolNumber, ops *tcpip.SocketOptions) {
+func (e *Endpoint) Init(s *stack.Stack, netProto tcpip.NetworkProtocolNumber, transProto tcpip.TransportProtocolNumber, ops *tcpip.SocketOptions, waiterQueue *waiter.Queue) {
 	e.mu.Lock()
 	memberships := e.multicastMemberships
 	e.mu.Unlock()
@@ -120,10 +131,11 @@ func (e *Endpoint) Init(s *stack.Stack, netProto tcpip.NetworkProtocolNumber, tr
 	}
 
 	*e = Endpoint{
-		stack:      s,
-		ops:        ops,
-		netProto:   netProto,
-		transProto: transProto,
+		stack:       s,
+		ops:         ops,
+		netProto:    netProto,
+		transProto:  transProto,
+		waiterQueue: waiterQueue,
 
 		info: stack.TransportEndpointInfo{
 			NetProto:   netProto,
@@ -217,11 +229,10 @@ func (e *Endpoint) calculateTTL(route *stack.Route) uint8 {
 
 // WriteContext holds the context for a write.
 type WriteContext struct {
-	transProto tcpip.TransportProtocolNumber
-	route      *stack.Route
-	ttl        uint8
-	tos        uint8
-	owner      tcpip.PacketOwner
+	e     *Endpoint
+	route *stack.Route
+	ttl   uint8
+	tos   uint8
 }
 
 // Release releases held resources.
@@ -249,21 +260,94 @@ func (c *WriteContext) PacketInfo() WritePacketInfo {
 	}
 }
 
+// TryNewPacketBuffer returns a new packet buffer iff the endpoint's send buffer
+// is not full.
+//
+// If this method returns nil, the caller should wait for the endpoint to become
+// writable.
+func (c *WriteContext) TryNewPacketBuffer(reserveHdrBytes int, data buffer.VectorisedView) *stack.PacketBuffer {
+	e := c.e
+
+	e.sendBufferSizeInUseMu.Lock()
+	defer e.sendBufferSizeInUseMu.Unlock()
+
+	if !e.hasSendSpaceRLocked() {
+		return nil
+	}
+
+	// Note that we allow oversubscription - if there is any space at all in the
+	// send buffer, we accept the full packet which may be larger than the space
+	// available. This is because if the endpoint reports that it is writable,
+	// a write operation should succeed.
+	//
+	// This matches Linux behaviour:
+	// https://github.com/torvalds/linux/blob/38d741cb70b/include/net/sock.h#L2519
+	// https://github.com/torvalds/linux/blob/38d741cb70b/net/core/sock.c#L2588
+	pktSize := int64(reserveHdrBytes) + int64(data.Size())
+	e.sendBufferSizeInUse += pktSize
+
+	return stack.NewPacketBuffer(stack.PacketBufferOptions{
+		ReserveHeaderBytes: reserveHdrBytes,
+		Data:               data,
+		OnRelease: func() {
+			e.sendBufferSizeInUseMu.Lock()
+			if got := e.sendBufferSizeInUse; got < pktSize {
+				e.sendBufferSizeInUseMu.Unlock()
+				panic(fmt.Sprintf("e.sendBufferSizeInUse=(%d) < pktSize(=%d)", got, pktSize))
+			}
+			e.sendBufferSizeInUse -= pktSize
+			signal := e.hasSendSpaceRLocked()
+			e.sendBufferSizeInUseMu.Unlock()
+
+			// Let waiters know if we now have space in the send buffer.
+			if signal {
+				e.waiterQueue.Notify(waiter.WritableEvents)
+			}
+		},
+	})
+}
+
 // WritePacket attempts to write the packet.
 func (c *WriteContext) WritePacket(pkt *stack.PacketBuffer, headerIncluded bool) tcpip.Error {
-	pkt.Owner = c.owner
+	c.e.mu.RLock()
+	pkt.Owner = c.e.owner
+	c.e.mu.RUnlock()
 
 	if headerIncluded {
 		return c.route.WriteHeaderIncludedPacket(pkt)
 	}
 
 	return c.route.WritePacket(stack.NetworkHeaderParams{
-		Protocol: c.transProto,
+		Protocol: c.e.transProto,
 		TTL:      c.ttl,
 		TOS:      c.tos,
 	}, pkt)
 }
 
+// MaybeSignalWritable signals waiters with writable events if the send buffer
+// has space.
+func (e *Endpoint) MaybeSignalWritable() {
+	e.sendBufferSizeInUseMu.RLock()
+	signal := e.hasSendSpaceRLocked()
+	e.sendBufferSizeInUseMu.RUnlock()
+
+	if signal {
+		e.waiterQueue.Notify(waiter.WritableEvents)
+	}
+}
+
+// HasSendSpace returns whether or not the send buffer has space.
+func (e *Endpoint) HasSendSpace() bool {
+	e.sendBufferSizeInUseMu.RLock()
+	defer e.sendBufferSizeInUseMu.RUnlock()
+	return e.hasSendSpaceRLocked()
+}
+
+// +checklocksread:e.sendBufferSizeInUseMu
+func (e *Endpoint) hasSendSpaceRLocked() bool {
+	return e.ops.GetSendBufferSize() > e.sendBufferSizeInUse
+}
+
 // AcquireContextForWrite acquires a WriteContext.
 func (e *Endpoint) AcquireContextForWrite(opts tcpip.WriteOptions) (WriteContext, tcpip.Error) {
 	e.mu.RLock()
@@ -348,11 +432,10 @@ func (e *Endpoint) AcquireContextForWrite(opts tcpip.WriteOptions) (WriteContext
 	}
 
 	return WriteContext{
-		transProto: e.transProto,
-		route:      route,
-		ttl:        ttl,
-		tos:        tos,
-		owner:      e.owner,
+		e:     e,
+		route: route,
+		ttl:   ttl,
+		tos:   tos,
 	}, nil
 }
 

pkg/tcpip/transport/internal/network/endpoint_test.go

@@ -36,6 +36,7 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip/transport"
 	"gvisor.dev/gvisor/pkg/tcpip/transport/internal/network"
 	"gvisor.dev/gvisor/pkg/tcpip/transport/udp"
+	"gvisor.dev/gvisor/pkg/waiter"
 )
 
 var (
@@ -150,7 +151,8 @@ func TestEndpointStateTransitions(t *testing.T) {
 
 			var ops tcpip.SocketOptions
 			var ep network.Endpoint
-			ep.Init(s, test.netProto, udp.ProtocolNumber, &ops)
+			var wq waiter.Queue
+			ep.Init(s, test.netProto, udp.ProtocolNumber, &ops, &wq)
 			defer ep.Close()
 			if state := ep.State(); state != transport.DatagramEndpointStateInitial {
 				t.Fatalf("got ep.State() = %s, want = %s", state, transport.DatagramEndpointStateInitial)
@@ -289,7 +291,8 @@ func TestBindNICID(t *testing.T) {
 
 					var ops tcpip.SocketOptions
 					var ep network.Endpoint
-					ep.Init(s, test.netProto, udp.ProtocolNumber, &ops)
+					var wq waiter.Queue
+					ep.Init(s, test.netProto, udp.ProtocolNumber, &ops, &wq)
 					defer ep.Close()
 					if ep.WasBound() {
 						t.Fatal("got ep.WasBound() = true, want = false")

pkg/tcpip/transport/raw/endpoint.go

@@ -133,7 +133,7 @@ func newEndpoint(s *stack.Stack, netProto tcpip.NetworkProtocolNumber, transProt
 	e.ops.SetHeaderIncluded(!associated)
 	e.ops.SetSendBufferSize(32*1024, false /* notify */)
 	e.ops.SetReceiveBufferSize(32*1024, false /* notify */)
-	e.net.Init(s, netProto, transProto, &e.ops)
+	e.net.Init(s, netProto, transProto, &e.ops, waiterQueue)
 
 	// Override with stack defaults.
 	var ss tcpip.SendBufferSizeOption
@@ -162,6 +162,11 @@ func newEndpoint(s *stack.Stack, netProto tcpip.NetworkProtocolNumber, transProt
 	return e, nil
 }
 
+// WakeupWriters implements tcpip.SocketOptionsHandler.
+func (e *endpoint) WakeupWriters() {
+	e.net.MaybeSignalWritable()
+}
+
 // HasNIC implements tcpip.SocketOptionsHandler.
 func (e *endpoint) HasNIC(id int32) bool {
 	return e.stack.HasNIC(tcpip.NICID(id))
@@ -353,10 +358,10 @@ func (e *endpoint) write(p tcpip.Payloader, opts tcpip.WriteOptions) (int64, tcp
 		header.PutChecksum(payloadBytes[ipv6ChecksumOffset:], ^xsum)
 	}
 
-	pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
-		ReserveHeaderBytes: int(ctx.PacketInfo().MaxHeaderLength),
-		Data:               buffer.View(payloadBytes).ToVectorisedView(),
-	})
+	pkt := ctx.TryNewPacketBuffer(int(ctx.PacketInfo().MaxHeaderLength), buffer.View(payloadBytes).ToVectorisedView())
+	if pkt == nil {
+		return 0, &tcpip.ErrWouldBlock{}
+	}
 	defer pkt.DecRef()
 
 	if err := ctx.WritePacket(pkt, e.ops.GetHeaderIncluded()); err != nil {
@@ -443,8 +448,11 @@ func (*endpoint) GetRemoteAddress() (tcpip.FullAddress, tcpip.Error) {
 
 // Readiness implements tcpip.Endpoint.Readiness.
 func (e *endpoint) Readiness(mask waiter.EventMask) waiter.EventMask {
-	// The endpoint is always writable.
-	result := waiter.WritableEvents & mask
+	var result waiter.EventMask
+
+	if e.net.HasSendSpace() {
+		result |= waiter.WritableEvents & mask
+	}
 
 	// Determine whether the endpoint is readable.
 	if (mask & waiter.ReadableEvents) != 0 {

pkg/tcpip/transport/udp/endpoint.go

@@ -115,7 +115,7 @@ func newEndpoint(s *stack.Stack, netProto tcpip.NetworkProtocolNumber, waiterQue
 	e.ops.SetMulticastLoop(true)
 	e.ops.SetSendBufferSize(32*1024, false /* notify */)
 	e.ops.SetReceiveBufferSize(32*1024, false /* notify */)
-	e.net.Init(s, netProto, header.UDPProtocolNumber, &e.ops)
+	e.net.Init(s, netProto, header.UDPProtocolNumber, &e.ops, waiterQueue)
 
 	// Override with stack defaults.
 	var ss tcpip.SendBufferSizeOption
@@ -131,6 +131,11 @@ func newEndpoint(s *stack.Stack, netProto tcpip.NetworkProtocolNumber, waiterQue
 	return e
 }
 
+// WakeupWriters implements tcpip.SocketOptionsHandler.
+func (e *endpoint) WakeupWriters() {
+	e.net.MaybeSignalWritable()
+}
+
 // UniqueID implements stack.TransportEndpoint.
 func (e *endpoint) UniqueID() uint64 {
 	return e.uniqueID
@@ -453,10 +458,10 @@ func (e *endpoint) write(p tcpip.Payloader, opts tcpip.WriteOptions) (int64, tcp
 	defer udpInfo.ctx.Release()
 
 	pktInfo := udpInfo.ctx.PacketInfo()
-	pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
-		ReserveHeaderBytes: header.UDPMinimumSize + int(pktInfo.MaxHeaderLength),
-		Data:               udpInfo.data.ToVectorisedView(),
-	})
+	pkt := udpInfo.ctx.TryNewPacketBuffer(header.UDPMinimumSize+int(pktInfo.MaxHeaderLength), udpInfo.data.ToVectorisedView())
+	if pkt == nil {
+		return 0, &tcpip.ErrWouldBlock{}
+	}
 	defer pkt.DecRef()
 
 	// Initialize the UDP header.
@@ -857,8 +862,11 @@ func (e *endpoint) GetRemoteAddress() (tcpip.FullAddress, tcpip.Error) {
 // Readiness returns the current readiness of the endpoint. For example, if
 // waiter.EventIn is set, the endpoint is immediately readable.
 func (e *endpoint) Readiness(mask waiter.EventMask) waiter.EventMask {
-	// The endpoint is always writable.
-	result := waiter.WritableEvents & mask
+	var result waiter.EventMask
+
+	if e.net.HasSendSpace() {
+		result |= waiter.WritableEvents & mask
+	}
 
 	// Determine if the endpoint is readable if requested.
 	if mask&waiter.ReadableEvents != 0 {

test/syscalls/linux/udp_socket.cc

@@ -2393,6 +2393,47 @@ TEST_P(UdpSocketControlMessagesTest, SetAndReceivePktInfo) {
   }
 }
 
+TEST_P(UdpSocketTest, SendPacketLargerThanSendBufOnNonBlockingSocket) {
+  constexpr int kSendBufSize = 4096;
+  ASSERT_THAT(setsockopt(sock_.get(), SOL_SOCKET, SO_SNDBUF, &kSendBufSize,
+                         sizeof(kSendBufSize)),
+              SyscallSucceeds());
+
+  // Set sock to non-blocking.
+  {
+    int opts = 0;
+    ASSERT_THAT(opts = fcntl(sock_.get(), F_GETFL), SyscallSucceeds());
+    ASSERT_THAT(fcntl(sock_.get(), F_SETFL, opts | O_NONBLOCK),
+                SyscallSucceeds());
+  }
+
+  {
+    sockaddr_storage addr = InetLoopbackAddr();
+    ASSERT_NO_ERRNO(BindSocket(sock_.get(), AsSockAddr(&addr)));
+  }
+
+  sockaddr_storage addr;
+  socklen_t len = sizeof(sockaddr_storage);
+  ASSERT_THAT(getsockname(sock_.get(), AsSockAddr(&addr), &len),
+              SyscallSucceeds());
+  ASSERT_EQ(len, addrlen_);
+
+  // We are allowed to send packets as large as we want as long as there is
+  // space in the send buffer, even if the new packet will result in more bytes
+  // being used than available in the send buffer.
+  char buf[kSendBufSize + 1];
+  ASSERT_THAT(
+      sendto(sock_.get(), buf, sizeof(buf), 0, AsSockAddr(&addr), sizeof(addr)),
+      SyscallSucceedsWithValue(sizeof(buf)));
+
+  // The second write may fail with EAGAIN if the previous send is still
+  // in-flight.
+  ASSERT_THAT(
+      sendto(sock_.get(), buf, sizeof(buf), 0, AsSockAddr(&addr), sizeof(addr)),
+      AnyOf(SyscallSucceedsWithValue(sizeof(buf)),
+            SyscallFailsWithErrno(EAGAIN)));
+}
+
 INSTANTIATE_TEST_SUITE_P(AllInetTests, UdpSocketControlMessagesTest,
                          ::testing::Values(AddressFamily::kIpv4,
                                            AddressFamily::kIpv6,