Initialize the send buffer handler in endpoint creation.

- This CL will initialize the function handler used for getting the send
buffer size limits during endpoint creation and does not require the caller of
SetSendBufferSize(..) to know the endpoint type(tcp/udp/..)

PiperOrigin-RevId: 353992634

16 files changed, 42 insertions, 38 deletions

diff --git a/pkg/sentry/socket/netstack/netstack.go b/pkg/sentry/socket/netstack/netstack.go
index b4d0651b8..7065a0e46 100644
--- a/pkg/sentry/socket/netstack/netstack.go
+++ b/pkg/sentry/socket/netstack/netstack.go
@@ -1615,20 +1615,15 @@ func setSockOptSocket(t *kernel.Task, s socket.SocketOps, ep commonEndpoint, nam
 			return syserr.ErrInvalidArgument
 		}
 
-		family, skType, skProto := s.Type()
+		family, _, _ := s.Type()
 		// TODO(gvisor.dev/issue/5132): We currently do not support
 		// setting this option for unix sockets.
 		if family == linux.AF_UNIX {
 			return nil
 		}
 
-		getBufferLimits := tcpip.GetStackSendBufferLimits
-		if isTCPSocket(skType, skProto) {
-			getBufferLimits = tcp.GetTCPSendBufferLimits
-		}
-
 		v := usermem.ByteOrder.Uint32(optVal)
-		ep.SocketOptions().SetSendBufferSize(int64(v), true, getBufferLimits)
+		ep.SocketOptions().SetSendBufferSize(int64(v), true)
 		return nil
 
 	case linux.SO_RCVBUF:
diff --git a/pkg/sentry/socket/unix/transport/connectioned.go b/pkg/sentry/socket/unix/transport/connectioned.go
index dbb7f7c31..b011082dc 100644
--- a/pkg/sentry/socket/unix/transport/connectioned.go
+++ b/pkg/sentry/socket/unix/transport/connectioned.go
@@ -128,7 +128,7 @@ func newConnectioned(ctx context.Context, stype linux.SockType, uid UniqueIDProv
 		idGenerator:  uid,
 		stype:        stype,
 	}
-	ep.ops.InitHandler(ep, nil)
+	ep.ops.InitHandler(ep, nil, nil)
 	return ep
 }
 
@@ -173,7 +173,7 @@ func NewExternal(ctx context.Context, stype linux.SockType, uid UniqueIDProvider
 		idGenerator:  uid,
 		stype:        stype,
 	}
-	ep.ops.InitHandler(ep, nil)
+	ep.ops.InitHandler(ep, nil, nil)
 	return ep
 }
 
@@ -296,7 +296,7 @@ func (e *connectionedEndpoint) BidirectionalConnect(ctx context.Context, ce Conn
 		idGenerator: e.idGenerator,
 		stype:       e.stype,
 	}
-	ne.ops.InitHandler(ne, nil)
+	ne.ops.InitHandler(ne, nil, nil)
 
 	readQueue := &queue{ReaderQueue: ce.WaiterQueue(), WriterQueue: ne.Queue, limit: initialLimit}
 	readQueue.InitRefs()
diff --git a/pkg/sentry/socket/unix/transport/connectionless.go b/pkg/sentry/socket/unix/transport/connectionless.go
index 895d2322e..20fa8b874 100644
--- a/pkg/sentry/socket/unix/transport/connectionless.go
+++ b/pkg/sentry/socket/unix/transport/connectionless.go
@@ -44,7 +44,7 @@ func NewConnectionless(ctx context.Context) Endpoint {
 	q := queue{ReaderQueue: ep.Queue, WriterQueue: &waiter.Queue{}, limit: initialLimit}
 	q.InitRefs()
 	ep.receiver = &queueReceiver{readQueue: &q}
-	ep.ops.InitHandler(ep, nil)
+	ep.ops.InitHandler(ep, nil, nil)
 	return ep
 }
 
diff --git a/pkg/tcpip/socketops.go b/pkg/tcpip/socketops.go
index 4149b91da..7eabbc599 100644
--- a/pkg/tcpip/socketops.go
+++ b/pkg/tcpip/socketops.go
@@ -205,6 +205,11 @@ type SocketOptions struct {
 	// bindToDevice determines the device to which the socket is bound.
 	bindToDevice int32
 
+	// getSendBufferLimits provides the handler to get the min, default and
+	// max size for send buffer. It  is initialized at the creation time and
+	// will not change.
+	getSendBufferLimits GetSendBufferLimits `state:"manual"`
+
 	// sendBufferSize determines the send buffer size for this socket.
 	sendBufferSize int64
 
@@ -218,9 +223,10 @@ type SocketOptions struct {
 
 // InitHandler initializes the handler. This must be called before using the
 // socket options utility.
-func (so *SocketOptions) InitHandler(handler SocketOptionsHandler, stack StackHandler) {
+func (so *SocketOptions) InitHandler(handler SocketOptionsHandler, stack StackHandler, getSendBufferLimits GetSendBufferLimits) {
 	so.handler = handler
 	so.stackHandler = stack
+	so.getSendBufferLimits = getSendBufferLimits
 }
 
 func storeAtomicBool(addr *uint32, v bool) {
@@ -568,14 +574,17 @@ func (so *SocketOptions) GetSendBufferSize() (int64, *Error) {
 
 // SetSendBufferSize sets value for SO_SNDBUF option. notify indicates if the
 // stack handler should be invoked to set the send buffer size.
-func (so *SocketOptions) SetSendBufferSize(sendBufferSize int64, notify bool, getBufferLimits GetSendBufferLimits) {
-	v := sendBufferSize
-	ss := getBufferLimits(so.stackHandler)
+func (so *SocketOptions) SetSendBufferSize(sendBufferSize int64, notify bool) {
+	if so.handler.IsUnixSocket() {
+		return
+	}
 
+	v := sendBufferSize
 	if notify {
 		// TODO(b/176170271): Notify waiters after size has grown.
 		// Make sure the send buffer size is within the min and max
 		// allowed.
+		ss := so.getSendBufferLimits(so.stackHandler)
 		min := int64(ss.Min)
 		max := int64(ss.Max)
 		// Validate the send buffer size with min and max values.
diff --git a/pkg/tcpip/stack/transport_test.go b/pkg/tcpip/stack/transport_test.go
index 8b86591c2..c49427c4c 100644
--- a/pkg/tcpip/stack/transport_test.go
+++ b/pkg/tcpip/stack/transport_test.go
@@ -70,7 +70,7 @@ func (f *fakeTransportEndpoint) SocketOptions() *tcpip.SocketOptions {
 
 func newFakeTransportEndpoint(proto *fakeTransportProtocol, netProto tcpip.NetworkProtocolNumber, s *stack.Stack) tcpip.Endpoint {
 	ep := &fakeTransportEndpoint{TransportEndpointInfo: stack.TransportEndpointInfo{NetProto: netProto}, proto: proto, uniqueID: s.UniqueID()}
-	ep.ops.InitHandler(ep, s)
+	ep.ops.InitHandler(ep, s, tcpip.GetStackSendBufferLimits)
 	return ep
 }
 
@@ -233,7 +233,7 @@ func (f *fakeTransportEndpoint) HandlePacket(id stack.TransportEndpointID, pkt *
 		peerAddr: route.RemoteAddress,
 		route:    route,
 	}
-	ep.ops.InitHandler(ep, f.proto.stack)
+	ep.ops.InitHandler(ep, f.proto.stack, tcpip.GetStackSendBufferLimits)
 	f.acceptQueue = append(f.acceptQueue, ep)
 }
 
diff --git a/pkg/tcpip/transport/icmp/endpoint.go b/pkg/tcpip/transport/icmp/endpoint.go
index c3d1252b1..e4bcd3120 100644
--- a/pkg/tcpip/transport/icmp/endpoint.go
+++ b/pkg/tcpip/transport/icmp/endpoint.go
@@ -96,13 +96,13 @@ func newEndpoint(s *stack.Stack, netProto tcpip.NetworkProtocolNumber, transProt
 		state:         stateInitial,
 		uniqueID:      s.UniqueID(),
 	}
-	ep.ops.InitHandler(ep, ep.stack)
-	ep.ops.SetSendBufferSize(32*1024, false /* notify */, tcpip.GetStackSendBufferLimits)
+	ep.ops.InitHandler(ep, ep.stack, tcpip.GetStackSendBufferLimits)
+	ep.ops.SetSendBufferSize(32*1024, false /* notify */)
 
 	// Override with stack defaults.
 	var ss tcpip.SendBufferSizeOption
 	if err := s.Option(&ss); err == nil {
-		ep.ops.SetSendBufferSize(int64(ss.Default), false /* notify */, tcpip.GetStackSendBufferLimits)
+		ep.ops.SetSendBufferSize(int64(ss.Default), false /* notify */)
 	}
 	return ep, nil
 }
diff --git a/pkg/tcpip/transport/icmp/endpoint_state.go b/pkg/tcpip/transport/icmp/endpoint_state.go
index 9335cbc5a..afe96998a 100644
--- a/pkg/tcpip/transport/icmp/endpoint_state.go
+++ b/pkg/tcpip/transport/icmp/endpoint_state.go
@@ -69,7 +69,7 @@ func (e *endpoint) afterLoad() {
 // Resume implements tcpip.ResumableEndpoint.Resume.
 func (e *endpoint) Resume(s *stack.Stack) {
 	e.stack = s
-	e.ops.InitHandler(e, e.stack)
+	e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits)
 
 	if e.state != stateBound && e.state != stateConnected {
 		return
diff --git a/pkg/tcpip/transport/packet/endpoint.go b/pkg/tcpip/transport/packet/endpoint.go
index 9ca9c9780..d48877677 100644
--- a/pkg/tcpip/transport/packet/endpoint.go
+++ b/pkg/tcpip/transport/packet/endpoint.go
@@ -105,12 +105,12 @@ func NewEndpoint(s *stack.Stack, cooked bool, netProto tcpip.NetworkProtocolNumb
 		waiterQueue:   waiterQueue,
 		rcvBufSizeMax: 32 * 1024,
 	}
-	ep.ops.InitHandler(ep, ep.stack)
+	ep.ops.InitHandler(ep, ep.stack, tcpip.GetStackSendBufferLimits)
 
 	// Override with stack defaults.
 	var ss tcpip.SendBufferSizeOption
 	if err := s.Option(&ss); err == nil {
-		ep.ops.SetSendBufferSize(int64(ss.Default), false /* notify */, tcpip.GetStackSendBufferLimits)
+		ep.ops.SetSendBufferSize(int64(ss.Default), false /* notify */)
 	}
 
 	var rs stack.ReceiveBufferSizeOption
diff --git a/pkg/tcpip/transport/packet/endpoint_state.go b/pkg/tcpip/transport/packet/endpoint_state.go
index 61221be71..4d98fb051 100644
--- a/pkg/tcpip/transport/packet/endpoint_state.go
+++ b/pkg/tcpip/transport/packet/endpoint_state.go
@@ -64,7 +64,7 @@ func (ep *endpoint) loadRcvBufSizeMax(max int) {
 // afterLoad is invoked by stateify.
 func (ep *endpoint) afterLoad() {
 	ep.stack = stack.StackFromEnv
-	ep.ops.InitHandler(ep, ep.stack)
+	ep.ops.InitHandler(ep, ep.stack, tcpip.GetStackSendBufferLimits)
 
 	// TODO(gvisor.dev/173): Once bind is supported, choose the right NIC.
 	if err := ep.stack.RegisterPacketEndpoint(0, ep.netProto, ep); err != nil {
diff --git a/pkg/tcpip/transport/raw/endpoint.go b/pkg/tcpip/transport/raw/endpoint.go
index 30644e92d..6c6d45188 100644
--- a/pkg/tcpip/transport/raw/endpoint.go
+++ b/pkg/tcpip/transport/raw/endpoint.go
@@ -112,14 +112,14 @@ func newEndpoint(s *stack.Stack, netProto tcpip.NetworkProtocolNumber, transProt
 		rcvBufSizeMax: 32 * 1024,
 		associated:    associated,
 	}
-	e.ops.InitHandler(e, e.stack)
+	e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits)
 	e.ops.SetHeaderIncluded(!associated)
-	e.ops.SetSendBufferSize(32*1024, false /* notify */, tcpip.GetStackSendBufferLimits)
+	e.ops.SetSendBufferSize(32*1024, false /* notify */)
 
 	// Override with stack defaults.
 	var ss tcpip.SendBufferSizeOption
 	if err := s.Option(&ss); err == nil {
-		e.ops.SetSendBufferSize(int64(ss.Default), false /* notify */, tcpip.GetStackSendBufferLimits)
+		e.ops.SetSendBufferSize(int64(ss.Default), false /* notify */)
 	}
 
 	var rs stack.ReceiveBufferSizeOption
diff --git a/pkg/tcpip/transport/raw/endpoint_state.go b/pkg/tcpip/transport/raw/endpoint_state.go
index f55e918ac..65c64d99f 100644
--- a/pkg/tcpip/transport/raw/endpoint_state.go
+++ b/pkg/tcpip/transport/raw/endpoint_state.go
@@ -69,7 +69,7 @@ func (e *endpoint) afterLoad() {
 // Resume implements tcpip.ResumableEndpoint.Resume.
 func (e *endpoint) Resume(s *stack.Stack) {
 	e.stack = s
-	e.ops.InitHandler(e, e.stack)
+	e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits)
 
 	// If the endpoint is connected, re-connect.
 	if e.connected {
diff --git a/pkg/tcpip/transport/tcp/endpoint.go b/pkg/tcpip/transport/tcp/endpoint.go
index eb6faaa4d..b6bd6d455 100644
--- a/pkg/tcpip/transport/tcp/endpoint.go
+++ b/pkg/tcpip/transport/tcp/endpoint.go
@@ -880,14 +880,14 @@ func newEndpoint(s *stack.Stack, netProto tcpip.NetworkProtocolNumber, waiterQue
 		windowClamp:   DefaultReceiveBufferSize,
 		maxSynRetries: DefaultSynRetries,
 	}
-	e.ops.InitHandler(e, e.stack)
+	e.ops.InitHandler(e, e.stack, GetTCPSendBufferLimits)
 	e.ops.SetMulticastLoop(true)
 	e.ops.SetQuickAck(true)
-	e.ops.SetSendBufferSize(DefaultSendBufferSize, false /* notify */, GetTCPSendBufferLimits)
+	e.ops.SetSendBufferSize(DefaultSendBufferSize, false /* notify */)
 
 	var ss tcpip.TCPSendBufferSizeRangeOption
 	if err := s.TransportProtocolOption(ProtocolNumber, &ss); err == nil {
-		e.ops.SetSendBufferSize(int64(ss.Default), false /* notify */, GetTCPSendBufferLimits)
+		e.ops.SetSendBufferSize(int64(ss.Default), false /* notify */)
 	}
 
 	var rs tcpip.TCPReceiveBufferSizeRangeOption
diff --git a/pkg/tcpip/transport/tcp/endpoint_state.go b/pkg/tcpip/transport/tcp/endpoint_state.go
index fa8ffaee0..4a01c81b4 100644
--- a/pkg/tcpip/transport/tcp/endpoint_state.go
+++ b/pkg/tcpip/transport/tcp/endpoint_state.go
@@ -181,7 +181,7 @@ func (e *endpoint) afterLoad() {
 // Resume implements tcpip.ResumableEndpoint.Resume.
 func (e *endpoint) Resume(s *stack.Stack) {
 	e.stack = s
-	e.ops.InitHandler(e, e.stack)
+	e.ops.InitHandler(e, e.stack, GetTCPSendBufferLimits)
 	e.segmentQueue.thaw()
 	epState := e.origEndpointState
 	switch epState {
diff --git a/pkg/tcpip/transport/tcp/tcp_test.go b/pkg/tcpip/transport/tcp/tcp_test.go
index 754f6b92a..87ff2b909 100644
--- a/pkg/tcpip/transport/tcp/tcp_test.go
+++ b/pkg/tcpip/transport/tcp/tcp_test.go
@@ -4473,7 +4473,7 @@ func TestMinMaxBufferSizes(t *testing.T) {
 
 	checkRecvBufferSize(t, ep, 200)
 
-	ep.SocketOptions().SetSendBufferSize(149, true, tcp.GetTCPSendBufferLimits)
+	ep.SocketOptions().SetSendBufferSize(149, true)
 
 	checkSendBufferSize(t, ep, 300)
 
@@ -4485,7 +4485,7 @@ func TestMinMaxBufferSizes(t *testing.T) {
 	// Values above max are capped at max and then doubled.
 	checkRecvBufferSize(t, ep, tcp.DefaultReceiveBufferSize*20*2)
 
-	ep.SocketOptions().SetSendBufferSize(1+tcp.DefaultSendBufferSize*30, true, tcp.GetTCPSendBufferLimits)
+	ep.SocketOptions().SetSendBufferSize(1+tcp.DefaultSendBufferSize*30, true)
 
 	// Values above max are capped at max and then doubled.
 	checkSendBufferSize(t, ep, tcp.DefaultSendBufferSize*30*2)
diff --git a/pkg/tcpip/transport/udp/endpoint.go b/pkg/tcpip/transport/udp/endpoint.go
index a3e513ff7..4988ba29b 100644
--- a/pkg/tcpip/transport/udp/endpoint.go
+++ b/pkg/tcpip/transport/udp/endpoint.go
@@ -178,14 +178,14 @@ func newEndpoint(s *stack.Stack, netProto tcpip.NetworkProtocolNumber, waiterQue
 		state:                StateInitial,
 		uniqueID:             s.UniqueID(),
 	}
-	e.ops.InitHandler(e, e.stack)
+	e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits)
 	e.ops.SetMulticastLoop(true)
-	e.ops.SetSendBufferSize(32*1024, false /* notify */, tcpip.GetStackSendBufferLimits)
+	e.ops.SetSendBufferSize(32*1024, false /* notify */)
 
 	// Override with stack defaults.
 	var ss tcpip.SendBufferSizeOption
 	if err := s.Option(&ss); err == nil {
-		e.ops.SetSendBufferSize(int64(ss.Default), false /* notify */, tcpip.GetStackSendBufferLimits)
+		e.ops.SetSendBufferSize(int64(ss.Default), false /* notify */)
 	}
 
 	var rs stack.ReceiveBufferSizeOption
diff --git a/pkg/tcpip/transport/udp/endpoint_state.go b/pkg/tcpip/transport/udp/endpoint_state.go
index 089574f05..feb53b553 100644
--- a/pkg/tcpip/transport/udp/endpoint_state.go
+++ b/pkg/tcpip/transport/udp/endpoint_state.go
@@ -91,7 +91,7 @@ func (e *endpoint) Resume(s *stack.Stack) {
 	defer e.mu.Unlock()
 
 	e.stack = s
-	e.ops.InitHandler(e, e.stack)
+	e.ops.InitHandler(e, e.stack, tcpip.GetStackSendBufferLimits)
 
 	for m := range e.multicastMemberships {
 		if err := e.stack.JoinGroup(e.NetProto, m.nicID, m.multicastAddr); err != nil {