conn.go 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880
  1. package socket
  2. import (
  3. "context"
  4. "errors"
  5. "io"
  6. "os"
  7. "sync"
  8. "sync/atomic"
  9. "syscall"
  10. "time"
  11. "golang.org/x/sys/unix"
  12. )
  13. // Lock in an expected public interface for convenience.
  14. var _ interface {
  15. io.ReadWriteCloser
  16. syscall.Conn
  17. SetDeadline(t time.Time) error
  18. SetReadDeadline(t time.Time) error
  19. SetWriteDeadline(t time.Time) error
  20. } = &Conn{}
  21. // A Conn is a low-level network connection which integrates with Go's runtime
  22. // network poller to provide asynchronous I/O and deadline support.
  23. //
  24. // Many of a Conn's blocking methods support net.Conn deadlines as well as
  25. // cancelation via context. Note that passing a context with a deadline set will
  26. // override any of the previous deadlines set by calls to the SetDeadline family
  27. // of methods.
  28. type Conn struct {
  29. // Indicates whether or not Conn.Close has been called. Must be accessed
  30. // atomically. Atomics definitions must come first in the Conn struct.
  31. closed uint32
  32. // A unique name for the Conn which is also associated with derived file
  33. // descriptors such as those created by accept(2).
  34. name string
  35. // facts contains information we have determined about Conn to trigger
  36. // alternate behavior in certain functions.
  37. facts facts
  38. // Provides access to the underlying file registered with the runtime
  39. // network poller, and arbitrary raw I/O calls.
  40. fd *os.File
  41. rc syscall.RawConn
  42. }
  43. // facts contains facts about a Conn.
  44. type facts struct {
  45. // isStream reports whether this is a streaming descriptor, as opposed to a
  46. // packet-based descriptor like a UDP socket.
  47. isStream bool
  48. // zeroReadIsEOF reports Whether a zero byte read indicates EOF. This is
  49. // false for a message based socket connection.
  50. zeroReadIsEOF bool
  51. }
  52. // A Config contains options for a Conn.
  53. type Config struct {
  54. // NetNS specifies the Linux network namespace the Conn will operate in.
  55. // This option is unsupported on other operating systems.
  56. //
  57. // If set (non-zero), Conn will enter the specified network namespace and an
  58. // error will occur in Socket if the operation fails.
  59. //
  60. // If not set (zero), a best-effort attempt will be made to enter the
  61. // network namespace of the calling thread: this means that any changes made
  62. // to the calling thread's network namespace will also be reflected in Conn.
  63. // If this operation fails (due to lack of permissions or because network
  64. // namespaces are disabled by kernel configuration), Socket will not return
  65. // an error, and the Conn will operate in the default network namespace of
  66. // the process. This enables non-privileged use of Conn in applications
  67. // which do not require elevated privileges.
  68. //
  69. // Entering a network namespace is a privileged operation (root or
  70. // CAP_SYS_ADMIN are required), and most applications should leave this set
  71. // to 0.
  72. NetNS int
  73. }
  74. // High-level methods which provide convenience over raw system calls.
  75. // Close closes the underlying file descriptor for the Conn, which also causes
  76. // all in-flight I/O operations to immediately unblock and return errors. Any
  77. // subsequent uses of Conn will result in EBADF.
  78. func (c *Conn) Close() error {
  79. // The caller has expressed an intent to close the socket, so immediately
  80. // increment s.closed to force further calls to result in EBADF before also
  81. // closing the file descriptor to unblock any outstanding operations.
  82. //
  83. // Because other operations simply check for s.closed != 0, we will permit
  84. // double Close, which would increment s.closed beyond 1.
  85. if atomic.AddUint32(&c.closed, 1) != 1 {
  86. // Multiple Close calls.
  87. return nil
  88. }
  89. return os.NewSyscallError("close", c.fd.Close())
  90. }
  91. // CloseRead shuts down the reading side of the Conn. Most callers should just
  92. // use Close.
  93. func (c *Conn) CloseRead() error { return c.Shutdown(unix.SHUT_RD) }
  94. // CloseWrite shuts down the writing side of the Conn. Most callers should just
  95. // use Close.
  96. func (c *Conn) CloseWrite() error { return c.Shutdown(unix.SHUT_WR) }
  97. // Read reads directly from the underlying file descriptor.
  98. func (c *Conn) Read(b []byte) (int, error) { return c.fd.Read(b) }
  99. // ReadContext reads from the underlying file descriptor with added support for
  100. // context cancelation.
  101. func (c *Conn) ReadContext(ctx context.Context, b []byte) (int, error) {
  102. if c.facts.isStream && len(b) > maxRW {
  103. b = b[:maxRW]
  104. }
  105. n, err := readT(c, ctx, "read", func(fd int) (int, error) {
  106. return unix.Read(fd, b)
  107. })
  108. if n == 0 && err == nil && c.facts.zeroReadIsEOF {
  109. return 0, io.EOF
  110. }
  111. return n, os.NewSyscallError("read", err)
  112. }
  113. // Write writes directly to the underlying file descriptor.
  114. func (c *Conn) Write(b []byte) (int, error) { return c.fd.Write(b) }
  115. // WriteContext writes to the underlying file descriptor with added support for
  116. // context cancelation.
  117. func (c *Conn) WriteContext(ctx context.Context, b []byte) (int, error) {
  118. var (
  119. n, nn int
  120. err error
  121. )
  122. doErr := c.write(ctx, "write", func(fd int) error {
  123. max := len(b)
  124. if c.facts.isStream && max-nn > maxRW {
  125. max = nn + maxRW
  126. }
  127. n, err = unix.Write(fd, b[nn:max])
  128. if n > 0 {
  129. nn += n
  130. }
  131. if nn == len(b) {
  132. return err
  133. }
  134. if n == 0 && err == nil {
  135. err = io.ErrUnexpectedEOF
  136. return nil
  137. }
  138. return err
  139. })
  140. if doErr != nil {
  141. return 0, doErr
  142. }
  143. return nn, os.NewSyscallError("write", err)
  144. }
  145. // SetDeadline sets both the read and write deadlines associated with the Conn.
  146. func (c *Conn) SetDeadline(t time.Time) error { return c.fd.SetDeadline(t) }
  147. // SetReadDeadline sets the read deadline associated with the Conn.
  148. func (c *Conn) SetReadDeadline(t time.Time) error { return c.fd.SetReadDeadline(t) }
  149. // SetWriteDeadline sets the write deadline associated with the Conn.
  150. func (c *Conn) SetWriteDeadline(t time.Time) error { return c.fd.SetWriteDeadline(t) }
  151. // ReadBuffer gets the size of the operating system's receive buffer associated
  152. // with the Conn.
  153. func (c *Conn) ReadBuffer() (int, error) {
  154. return c.GetsockoptInt(unix.SOL_SOCKET, unix.SO_RCVBUF)
  155. }
  156. // WriteBuffer gets the size of the operating system's transmit buffer
  157. // associated with the Conn.
  158. func (c *Conn) WriteBuffer() (int, error) {
  159. return c.GetsockoptInt(unix.SOL_SOCKET, unix.SO_SNDBUF)
  160. }
  161. // SetReadBuffer sets the size of the operating system's receive buffer
  162. // associated with the Conn.
  163. //
  164. // When called with elevated privileges on Linux, the SO_RCVBUFFORCE option will
  165. // be used to override operating system limits. Otherwise SO_RCVBUF is used
  166. // (which obeys operating system limits).
  167. func (c *Conn) SetReadBuffer(bytes int) error { return c.setReadBuffer(bytes) }
  168. // SetWriteBuffer sets the size of the operating system's transmit buffer
  169. // associated with the Conn.
  170. //
  171. // When called with elevated privileges on Linux, the SO_SNDBUFFORCE option will
  172. // be used to override operating system limits. Otherwise SO_SNDBUF is used
  173. // (which obeys operating system limits).
  174. func (c *Conn) SetWriteBuffer(bytes int) error { return c.setWriteBuffer(bytes) }
  175. // SyscallConn returns a raw network connection. This implements the
  176. // syscall.Conn interface.
  177. //
  178. // SyscallConn is intended for advanced use cases, such as getting and setting
  179. // arbitrary socket options using the socket's file descriptor. If possible,
  180. // those operations should be performed using methods on Conn instead.
  181. //
  182. // Once invoked, it is the caller's responsibility to ensure that operations
  183. // performed using Conn and the syscall.RawConn do not conflict with each other.
  184. func (c *Conn) SyscallConn() (syscall.RawConn, error) {
  185. if atomic.LoadUint32(&c.closed) != 0 {
  186. return nil, os.NewSyscallError("syscallconn", unix.EBADF)
  187. }
  188. // TODO(mdlayher): mutex or similar to enforce syscall.RawConn contract of
  189. // FD remaining valid for duration of calls?
  190. return c.rc, nil
  191. }
  192. // Socket wraps the socket(2) system call to produce a Conn. domain, typ, and
  193. // proto are passed directly to socket(2), and name should be a unique name for
  194. // the socket type such as "netlink" or "vsock".
  195. //
  196. // The cfg parameter specifies optional configuration for the Conn. If nil, no
  197. // additional configuration will be applied.
  198. //
  199. // If the operating system supports SOCK_CLOEXEC and SOCK_NONBLOCK, they are
  200. // automatically applied to typ to mirror the standard library's socket flag
  201. // behaviors.
  202. func Socket(domain, typ, proto int, name string, cfg *Config) (*Conn, error) {
  203. if cfg == nil {
  204. cfg = &Config{}
  205. }
  206. if cfg.NetNS == 0 {
  207. // Non-Linux or no network namespace.
  208. return socket(domain, typ, proto, name)
  209. }
  210. // Linux only: create Conn in the specified network namespace.
  211. return withNetNS(cfg.NetNS, func() (*Conn, error) {
  212. return socket(domain, typ, proto, name)
  213. })
  214. }
  215. // socket is the internal, cross-platform entry point for socket(2).
  216. func socket(domain, typ, proto int, name string) (*Conn, error) {
  217. var (
  218. fd int
  219. err error
  220. )
  221. for {
  222. fd, err = unix.Socket(domain, typ|socketFlags, proto)
  223. switch {
  224. case err == nil:
  225. // Some OSes already set CLOEXEC with typ.
  226. if !flagCLOEXEC {
  227. unix.CloseOnExec(fd)
  228. }
  229. // No error, prepare the Conn.
  230. return New(fd, name)
  231. case !ready(err):
  232. // System call interrupted or not ready, try again.
  233. continue
  234. case err == unix.EINVAL, err == unix.EPROTONOSUPPORT:
  235. // On Linux, SOCK_NONBLOCK and SOCK_CLOEXEC were introduced in
  236. // 2.6.27. On FreeBSD, both flags were introduced in FreeBSD 10.
  237. // EINVAL and EPROTONOSUPPORT check for earlier versions of these
  238. // OSes respectively.
  239. //
  240. // Mirror what the standard library does when creating file
  241. // descriptors: avoid racing a fork/exec with the creation of new
  242. // file descriptors, so that child processes do not inherit socket
  243. // file descriptors unexpectedly.
  244. //
  245. // For a more thorough explanation, see similar work in the Go tree:
  246. // func sysSocket in net/sock_cloexec.go, as well as the detailed
  247. // comment in syscall/exec_unix.go.
  248. syscall.ForkLock.RLock()
  249. fd, err = unix.Socket(domain, typ, proto)
  250. if err != nil {
  251. syscall.ForkLock.RUnlock()
  252. return nil, os.NewSyscallError("socket", err)
  253. }
  254. unix.CloseOnExec(fd)
  255. syscall.ForkLock.RUnlock()
  256. return New(fd, name)
  257. default:
  258. // Unhandled error.
  259. return nil, os.NewSyscallError("socket", err)
  260. }
  261. }
  262. }
  263. // FileConn returns a copy of the network connection corresponding to the open
  264. // file. It is the caller's responsibility to close the file when finished.
  265. // Closing the Conn does not affect the File, and closing the File does not
  266. // affect the Conn.
  267. func FileConn(f *os.File, name string) (*Conn, error) {
  268. // First we'll try to do fctnl(2) with F_DUPFD_CLOEXEC because we can dup
  269. // the file descriptor and set the flag in one syscall.
  270. fd, err := unix.FcntlInt(f.Fd(), unix.F_DUPFD_CLOEXEC, 0)
  271. switch err {
  272. case nil:
  273. // OK, ready to set up non-blocking I/O.
  274. return New(fd, name)
  275. case unix.EINVAL:
  276. // The kernel rejected our fcntl(2), fall back to separate dup(2) and
  277. // setting close on exec.
  278. //
  279. // Mirror what the standard library does when creating file descriptors:
  280. // avoid racing a fork/exec with the creation of new file descriptors,
  281. // so that child processes do not inherit socket file descriptors
  282. // unexpectedly.
  283. syscall.ForkLock.RLock()
  284. fd, err := unix.Dup(fd)
  285. if err != nil {
  286. syscall.ForkLock.RUnlock()
  287. return nil, os.NewSyscallError("dup", err)
  288. }
  289. unix.CloseOnExec(fd)
  290. syscall.ForkLock.RUnlock()
  291. return New(fd, name)
  292. default:
  293. // Any other errors.
  294. return nil, os.NewSyscallError("fcntl", err)
  295. }
  296. }
  297. // New wraps an existing file descriptor to create a Conn. name should be a
  298. // unique name for the socket type such as "netlink" or "vsock".
  299. //
  300. // Most callers should use Socket or FileConn to construct a Conn. New is
  301. // intended for integrating with specific system calls which provide a file
  302. // descriptor that supports asynchronous I/O. The file descriptor is immediately
  303. // set to nonblocking mode and registered with Go's runtime network poller for
  304. // future I/O operations.
  305. //
  306. // Unlike FileConn, New does not duplicate the existing file descriptor in any
  307. // way. The returned Conn takes ownership of the underlying file descriptor.
  308. func New(fd int, name string) (*Conn, error) {
  309. // All Conn I/O is nonblocking for integration with Go's runtime network
  310. // poller. Depending on the OS this might already be set but it can't hurt
  311. // to set it again.
  312. if err := unix.SetNonblock(fd, true); err != nil {
  313. return nil, os.NewSyscallError("setnonblock", err)
  314. }
  315. // os.NewFile registers the non-blocking file descriptor with the runtime
  316. // poller, which is then used for most subsequent operations except those
  317. // that require raw I/O via SyscallConn.
  318. //
  319. // See also: https://golang.org/pkg/os/#NewFile
  320. f := os.NewFile(uintptr(fd), name)
  321. rc, err := f.SyscallConn()
  322. if err != nil {
  323. return nil, err
  324. }
  325. c := &Conn{
  326. name: name,
  327. fd: f,
  328. rc: rc,
  329. }
  330. // Probe the file descriptor for socket settings.
  331. sotype, err := c.GetsockoptInt(unix.SOL_SOCKET, unix.SO_TYPE)
  332. switch {
  333. case err == nil:
  334. // File is a socket, check its properties.
  335. c.facts = facts{
  336. isStream: sotype == unix.SOCK_STREAM,
  337. zeroReadIsEOF: sotype != unix.SOCK_DGRAM && sotype != unix.SOCK_RAW,
  338. }
  339. case errors.Is(err, unix.ENOTSOCK):
  340. // File is not a socket, treat it as a regular file.
  341. c.facts = facts{
  342. isStream: true,
  343. zeroReadIsEOF: true,
  344. }
  345. default:
  346. return nil, err
  347. }
  348. return c, nil
  349. }
  350. // Low-level methods which provide raw system call access.
  351. // Accept wraps accept(2) or accept4(2) depending on the operating system, but
  352. // returns a Conn for the accepted connection rather than a raw file descriptor.
  353. //
  354. // If the operating system supports accept4(2) (which allows flags),
  355. // SOCK_CLOEXEC and SOCK_NONBLOCK are automatically applied to flags to mirror
  356. // the standard library's socket flag behaviors.
  357. //
  358. // If the operating system only supports accept(2) (which does not allow flags)
  359. // and flags is not zero, an error will be returned.
  360. //
  361. // Accept obeys context cancelation and uses the deadline set on the context to
  362. // cancel accepting the next connection. If a deadline is set on ctx, this
  363. // deadline will override any previous deadlines set using SetDeadline or
  364. // SetReadDeadline. Upon return, the read deadline is cleared.
  365. func (c *Conn) Accept(ctx context.Context, flags int) (*Conn, unix.Sockaddr, error) {
  366. type ret struct {
  367. nfd int
  368. sa unix.Sockaddr
  369. }
  370. r, err := readT(c, ctx, sysAccept, func(fd int) (ret, error) {
  371. // Either accept(2) or accept4(2) depending on the OS.
  372. nfd, sa, err := accept(fd, flags|socketFlags)
  373. return ret{nfd, sa}, err
  374. })
  375. if err != nil {
  376. // internal/poll, context error, or user function error.
  377. return nil, nil, err
  378. }
  379. // Successfully accepted a connection, wrap it in a Conn for use by the
  380. // caller.
  381. ac, err := New(r.nfd, c.name)
  382. if err != nil {
  383. return nil, nil, err
  384. }
  385. return ac, r.sa, nil
  386. }
  387. // Bind wraps bind(2).
  388. func (c *Conn) Bind(sa unix.Sockaddr) error {
  389. return c.control(context.Background(), "bind", func(fd int) error {
  390. return unix.Bind(fd, sa)
  391. })
  392. }
  393. // Connect wraps connect(2). In order to verify that the underlying socket is
  394. // connected to a remote peer, Connect calls getpeername(2) and returns the
  395. // unix.Sockaddr from that call.
  396. //
  397. // Connect obeys context cancelation and uses the deadline set on the context to
  398. // cancel connecting to a remote peer. If a deadline is set on ctx, this
  399. // deadline will override any previous deadlines set using SetDeadline or
  400. // SetWriteDeadline. Upon return, the write deadline is cleared.
  401. func (c *Conn) Connect(ctx context.Context, sa unix.Sockaddr) (unix.Sockaddr, error) {
  402. const op = "connect"
  403. // TODO(mdlayher): it would seem that trying to connect to unbound vsock
  404. // listeners by calling Connect multiple times results in ECONNRESET for the
  405. // first and nil error for subsequent calls. Do we need to memoize the
  406. // error? Check what the stdlib behavior is.
  407. var (
  408. // Track progress between invocations of the write closure. We don't
  409. // have an explicit WaitWrite call like internal/poll does, so we have
  410. // to wait until the runtime calls the closure again to indicate we can
  411. // write.
  412. progress uint32
  413. // Capture closure sockaddr and error.
  414. rsa unix.Sockaddr
  415. err error
  416. )
  417. doErr := c.write(ctx, op, func(fd int) error {
  418. if atomic.AddUint32(&progress, 1) == 1 {
  419. // First call: initiate connect.
  420. return unix.Connect(fd, sa)
  421. }
  422. // Subsequent calls: the runtime network poller indicates fd is
  423. // writable. Check for errno.
  424. errno, gerr := c.GetsockoptInt(unix.SOL_SOCKET, unix.SO_ERROR)
  425. if gerr != nil {
  426. return gerr
  427. }
  428. if errno != 0 {
  429. // Connection is still not ready or failed. If errno indicates
  430. // the socket is not ready, we will wait for the next write
  431. // event. Otherwise we propagate this errno back to the as a
  432. // permanent error.
  433. uerr := unix.Errno(errno)
  434. err = uerr
  435. return uerr
  436. }
  437. // According to internal/poll, it's possible for the runtime network
  438. // poller to spuriously wake us and return errno 0 for SO_ERROR.
  439. // Make sure we are actually connected to a peer.
  440. peer, err := c.Getpeername()
  441. if err != nil {
  442. // internal/poll unconditionally goes back to WaitWrite.
  443. // Synthesize an error that will do the same for us.
  444. return unix.EAGAIN
  445. }
  446. // Connection complete.
  447. rsa = peer
  448. return nil
  449. })
  450. if doErr != nil {
  451. // internal/poll or context error.
  452. return nil, doErr
  453. }
  454. if err == unix.EISCONN {
  455. // TODO(mdlayher): is this block obsolete with the addition of the
  456. // getsockopt SO_ERROR check above?
  457. //
  458. // EISCONN is reported if the socket is already established and should
  459. // not be treated as an error.
  460. // - Darwin reports this for at least TCP sockets
  461. // - Linux reports this for at least AF_VSOCK sockets
  462. return rsa, nil
  463. }
  464. return rsa, os.NewSyscallError(op, err)
  465. }
  466. // Getsockname wraps getsockname(2).
  467. func (c *Conn) Getsockname() (unix.Sockaddr, error) {
  468. return controlT(c, context.Background(), "getsockname", unix.Getsockname)
  469. }
  470. // Getpeername wraps getpeername(2).
  471. func (c *Conn) Getpeername() (unix.Sockaddr, error) {
  472. return controlT(c, context.Background(), "getpeername", unix.Getpeername)
  473. }
  474. // GetsockoptInt wraps getsockopt(2) for integer values.
  475. func (c *Conn) GetsockoptInt(level, opt int) (int, error) {
  476. return controlT(c, context.Background(), "getsockopt", func(fd int) (int, error) {
  477. return unix.GetsockoptInt(fd, level, opt)
  478. })
  479. }
  480. // Listen wraps listen(2).
  481. func (c *Conn) Listen(n int) error {
  482. return c.control(context.Background(), "listen", func(fd int) error {
  483. return unix.Listen(fd, n)
  484. })
  485. }
  486. // Recvmsg wraps recvmsg(2).
  487. func (c *Conn) Recvmsg(ctx context.Context, p, oob []byte, flags int) (int, int, int, unix.Sockaddr, error) {
  488. type ret struct {
  489. n, oobn, recvflags int
  490. from unix.Sockaddr
  491. }
  492. r, err := readT(c, ctx, "recvmsg", func(fd int) (ret, error) {
  493. n, oobn, recvflags, from, err := unix.Recvmsg(fd, p, oob, flags)
  494. return ret{n, oobn, recvflags, from}, err
  495. })
  496. if r.n == 0 && err == nil && c.facts.zeroReadIsEOF {
  497. return 0, 0, 0, nil, io.EOF
  498. }
  499. return r.n, r.oobn, r.recvflags, r.from, err
  500. }
  501. // Recvfrom wraps recvfrom(2).
  502. func (c *Conn) Recvfrom(ctx context.Context, p []byte, flags int) (int, unix.Sockaddr, error) {
  503. type ret struct {
  504. n int
  505. addr unix.Sockaddr
  506. }
  507. out, err := readT(c, ctx, "recvfrom", func(fd int) (ret, error) {
  508. n, addr, err := unix.Recvfrom(fd, p, flags)
  509. return ret{n, addr}, err
  510. })
  511. if out.n == 0 && err == nil && c.facts.zeroReadIsEOF {
  512. return 0, nil, io.EOF
  513. }
  514. return out.n, out.addr, err
  515. }
  516. // Sendmsg wraps sendmsg(2).
  517. func (c *Conn) Sendmsg(ctx context.Context, p, oob []byte, to unix.Sockaddr, flags int) (int, error) {
  518. return writeT(c, ctx, "sendmsg", func(fd int) (int, error) {
  519. return unix.SendmsgN(fd, p, oob, to, flags)
  520. })
  521. }
  522. // Sendto wraps sendto(2).
  523. func (c *Conn) Sendto(ctx context.Context, p []byte, flags int, to unix.Sockaddr) error {
  524. return c.write(ctx, "sendto", func(fd int) error {
  525. return unix.Sendto(fd, p, flags, to)
  526. })
  527. }
  528. // SetsockoptInt wraps setsockopt(2) for integer values.
  529. func (c *Conn) SetsockoptInt(level, opt, value int) error {
  530. return c.control(context.Background(), "setsockopt", func(fd int) error {
  531. return unix.SetsockoptInt(fd, level, opt, value)
  532. })
  533. }
  534. // Shutdown wraps shutdown(2).
  535. func (c *Conn) Shutdown(how int) error {
  536. return c.control(context.Background(), "shutdown", func(fd int) error {
  537. return unix.Shutdown(fd, how)
  538. })
  539. }
  540. // Conn low-level read/write/control functions. These functions mirror the
  541. // syscall.RawConn APIs but the input closures return errors rather than
  542. // booleans.
  543. // read wraps readT to execute a function and capture its error result. This is
  544. // a convenience wrapper for functions which don't return any extra values.
  545. func (c *Conn) read(ctx context.Context, op string, f func(fd int) error) error {
  546. _, err := readT(c, ctx, op, func(fd int) (struct{}, error) {
  547. return struct{}{}, f(fd)
  548. })
  549. return err
  550. }
  551. // write executes f, a write function, against the associated file descriptor.
  552. // op is used to create an *os.SyscallError if the file descriptor is closed.
  553. func (c *Conn) write(ctx context.Context, op string, f func(fd int) error) error {
  554. _, err := writeT(c, ctx, op, func(fd int) (struct{}, error) {
  555. return struct{}{}, f(fd)
  556. })
  557. return err
  558. }
  559. // readT executes c.rc.Read for op using the input function, returning a newly
  560. // allocated result T.
  561. func readT[T any](c *Conn, ctx context.Context, op string, f func(fd int) (T, error)) (T, error) {
  562. return rwT(c, rwContext[T]{
  563. Context: ctx,
  564. Type: read,
  565. Op: op,
  566. Do: f,
  567. })
  568. }
  569. // writeT executes c.rc.Write for op using the input function, returning a newly
  570. // allocated result T.
  571. func writeT[T any](c *Conn, ctx context.Context, op string, f func(fd int) (T, error)) (T, error) {
  572. return rwT(c, rwContext[T]{
  573. Context: ctx,
  574. Type: write,
  575. Op: op,
  576. Do: f,
  577. })
  578. }
  579. // readWrite indicates if an operation intends to read or write.
  580. type readWrite bool
  581. // Possible readWrite values.
  582. const (
  583. read readWrite = false
  584. write readWrite = true
  585. )
  586. // An rwContext provides arguments to rwT.
  587. type rwContext[T any] struct {
  588. // The caller's context passed for cancelation.
  589. Context context.Context
  590. // The type of an operation: read or write.
  591. Type readWrite
  592. // The name of the operation used in errors.
  593. Op string
  594. // The actual function to perform.
  595. Do func(fd int) (T, error)
  596. }
  597. // rwT executes c.rc.Read or c.rc.Write (depending on the value of rw.Type) for
  598. // rw.Op using the input function, returning a newly allocated result T.
  599. //
  600. // It obeys context cancelation and the rw.Context must not be nil.
  601. func rwT[T any](c *Conn, rw rwContext[T]) (T, error) {
  602. if atomic.LoadUint32(&c.closed) != 0 {
  603. // If the file descriptor is already closed, do nothing.
  604. return *new(T), os.NewSyscallError(rw.Op, unix.EBADF)
  605. }
  606. if err := rw.Context.Err(); err != nil {
  607. // Early exit due to context cancel.
  608. return *new(T), os.NewSyscallError(rw.Op, err)
  609. }
  610. var (
  611. // The read or write function used to access the runtime network poller.
  612. poll func(func(uintptr) bool) error
  613. // The read or write function used to set the matching deadline.
  614. deadline func(time.Time) error
  615. )
  616. if rw.Type == write {
  617. poll = c.rc.Write
  618. deadline = c.SetWriteDeadline
  619. } else {
  620. poll = c.rc.Read
  621. deadline = c.SetReadDeadline
  622. }
  623. var (
  624. // Whether or not the context carried a deadline we are actively using
  625. // for cancelation.
  626. setDeadline bool
  627. // Signals for the cancelation watcher goroutine.
  628. wg sync.WaitGroup
  629. doneC = make(chan struct{})
  630. // Atomic: reports whether we have to disarm the deadline.
  631. //
  632. // TODO(mdlayher): switch back to atomic.Bool when we drop support for
  633. // Go 1.18.
  634. needDisarm int64
  635. )
  636. // On cancel, clean up the watcher.
  637. defer func() {
  638. close(doneC)
  639. wg.Wait()
  640. }()
  641. if d, ok := rw.Context.Deadline(); ok {
  642. // The context has an explicit deadline. We will use it for cancelation
  643. // but disarm it after poll for the next call.
  644. if err := deadline(d); err != nil {
  645. return *new(T), err
  646. }
  647. setDeadline = true
  648. atomic.AddInt64(&needDisarm, 1)
  649. } else {
  650. // The context does not have an explicit deadline. We have to watch for
  651. // cancelation so we can propagate that signal to immediately unblock
  652. // the runtime network poller.
  653. //
  654. // TODO(mdlayher): is it possible to detect a background context vs a
  655. // context with possible future cancel?
  656. wg.Add(1)
  657. go func() {
  658. defer wg.Done()
  659. select {
  660. case <-rw.Context.Done():
  661. // Cancel the operation. Make the caller disarm after poll
  662. // returns.
  663. atomic.AddInt64(&needDisarm, 1)
  664. _ = deadline(time.Unix(0, 1))
  665. case <-doneC:
  666. // Nothing to do.
  667. }
  668. }()
  669. }
  670. var (
  671. t T
  672. err error
  673. )
  674. pollErr := poll(func(fd uintptr) bool {
  675. t, err = rw.Do(int(fd))
  676. return ready(err)
  677. })
  678. if atomic.LoadInt64(&needDisarm) > 0 {
  679. _ = deadline(time.Time{})
  680. }
  681. if pollErr != nil {
  682. if rw.Context.Err() != nil || (setDeadline && errors.Is(pollErr, os.ErrDeadlineExceeded)) {
  683. // The caller canceled the operation or we set a deadline internally
  684. // and it was reached.
  685. //
  686. // Unpack a plain context error. We wait for the context to be done
  687. // to synchronize state externally. Otherwise we have noticed I/O
  688. // timeout wakeups when we set a deadline but the context was not
  689. // yet marked done.
  690. <-rw.Context.Done()
  691. return *new(T), os.NewSyscallError(rw.Op, rw.Context.Err())
  692. }
  693. // Error from syscall.RawConn methods. Conventionally the standard
  694. // library does not wrap internal/poll errors in os.NewSyscallError.
  695. return *new(T), pollErr
  696. }
  697. // Result from user function.
  698. return t, os.NewSyscallError(rw.Op, err)
  699. }
  700. // control executes Conn.control for op using the input function.
  701. func (c *Conn) control(ctx context.Context, op string, f func(fd int) error) error {
  702. _, err := controlT(c, ctx, op, func(fd int) (struct{}, error) {
  703. return struct{}{}, f(fd)
  704. })
  705. return err
  706. }
  707. // controlT executes c.rc.Control for op using the input function, returning a
  708. // newly allocated result T.
  709. func controlT[T any](c *Conn, ctx context.Context, op string, f func(fd int) (T, error)) (T, error) {
  710. if atomic.LoadUint32(&c.closed) != 0 {
  711. // If the file descriptor is already closed, do nothing.
  712. return *new(T), os.NewSyscallError(op, unix.EBADF)
  713. }
  714. var (
  715. t T
  716. err error
  717. )
  718. doErr := c.rc.Control(func(fd uintptr) {
  719. // Repeatedly attempt the syscall(s) invoked by f until completion is
  720. // indicated by the return value of ready or the context is canceled.
  721. //
  722. // The last values for t and err are captured outside of the closure for
  723. // use when the loop breaks.
  724. for {
  725. if err = ctx.Err(); err != nil {
  726. // Early exit due to context cancel.
  727. return
  728. }
  729. t, err = f(int(fd))
  730. if ready(err) {
  731. return
  732. }
  733. }
  734. })
  735. if doErr != nil {
  736. // Error from syscall.RawConn methods. Conventionally the standard
  737. // library does not wrap internal/poll errors in os.NewSyscallError.
  738. return *new(T), doErr
  739. }
  740. // Result from user function.
  741. return t, os.NewSyscallError(op, err)
  742. }
  743. // ready indicates readiness based on the value of err.
  744. func ready(err error) bool {
  745. switch err {
  746. case unix.EAGAIN, unix.EINPROGRESS, unix.EINTR:
  747. // When a socket is in non-blocking mode, we might see a variety of errors:
  748. // - EAGAIN: most common case for a socket read not being ready
  749. // - EINPROGRESS: reported by some sockets when first calling connect
  750. // - EINTR: system call interrupted, more frequently occurs in Go 1.14+
  751. // because goroutines can be asynchronously preempted
  752. //
  753. // Return false to let the poller wait for readiness. See the source code
  754. // for internal/poll.FD.RawRead for more details.
  755. return false
  756. default:
  757. // Ready regardless of whether there was an error or no error.
  758. return true
  759. }
  760. }
  761. // Darwin and FreeBSD can't read or write 2GB+ files at a time,
  762. // even on 64-bit systems.
  763. // The same is true of socket implementations on many systems.
  764. // See golang.org/issue/7812 and golang.org/issue/16266.
  765. // Use 1GB instead of, say, 2GB-1, to keep subsequent reads aligned.
  766. const maxRW = 1 << 30