rwcancel.go raw

   1  //go:build !windows && !wasm
   2  
   3  /* SPDX-License-Identifier: MIT
   4   *
   5   * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved.
   6   */
   7  
   8  // Package rwcancel implements cancelable read/write operations on
   9  // a file descriptor.
  10  package rwcancel
  11  
  12  import (
  13  	"errors"
  14  	"os"
  15  	"syscall"
  16  
  17  	"golang.org/x/sys/unix"
  18  )
  19  
  20  type RWCancel struct {
  21  	fd            int
  22  	closingReader *os.File
  23  	closingWriter *os.File
  24  }
  25  
  26  func NewRWCancel(fd int) (*RWCancel, error) {
  27  	err := unix.SetNonblock(fd, true)
  28  	if err != nil {
  29  		return nil, err
  30  	}
  31  	rwcancel := RWCancel{fd: fd}
  32  
  33  	rwcancel.closingReader, rwcancel.closingWriter, err = os.Pipe()
  34  	if err != nil {
  35  		return nil, err
  36  	}
  37  
  38  	return &rwcancel, nil
  39  }
  40  
  41  func RetryAfterError(err error) bool {
  42  	return errors.Is(err, syscall.EAGAIN) || errors.Is(err, syscall.EINTR)
  43  }
  44  
  45  func (rw *RWCancel) ReadyRead() bool {
  46  	closeFd := int32(rw.closingReader.Fd())
  47  
  48  	pollFds := []unix.PollFd{{Fd: int32(rw.fd), Events: unix.POLLIN}, {Fd: closeFd, Events: unix.POLLIN}}
  49  	var err error
  50  	for {
  51  		_, err = unix.Poll(pollFds, -1)
  52  		if err == nil || !RetryAfterError(err) {
  53  			break
  54  		}
  55  	}
  56  	if err != nil {
  57  		return false
  58  	}
  59  	if pollFds[1].Revents != 0 {
  60  		return false
  61  	}
  62  	return pollFds[0].Revents != 0
  63  }
  64  
  65  func (rw *RWCancel) ReadyWrite() bool {
  66  	closeFd := int32(rw.closingReader.Fd())
  67  	pollFds := []unix.PollFd{{Fd: int32(rw.fd), Events: unix.POLLOUT}, {Fd: closeFd, Events: unix.POLLIN}}
  68  	var err error
  69  	for {
  70  		_, err = unix.Poll(pollFds, -1)
  71  		if err == nil || !RetryAfterError(err) {
  72  			break
  73  		}
  74  	}
  75  	if err != nil {
  76  		return false
  77  	}
  78  
  79  	if pollFds[1].Revents != 0 {
  80  		return false
  81  	}
  82  	return pollFds[0].Revents != 0
  83  }
  84  
  85  func (rw *RWCancel) Read(p []byte) (n int, err error) {
  86  	for {
  87  		n, err := unix.Read(rw.fd, p)
  88  		if err == nil || !RetryAfterError(err) {
  89  			return n, err
  90  		}
  91  		if !rw.ReadyRead() {
  92  			return 0, os.ErrClosed
  93  		}
  94  	}
  95  }
  96  
  97  func (rw *RWCancel) Write(p []byte) (n int, err error) {
  98  	for {
  99  		n, err := unix.Write(rw.fd, p)
 100  		if err == nil || !RetryAfterError(err) {
 101  			return n, err
 102  		}
 103  		if !rw.ReadyWrite() {
 104  			return 0, os.ErrClosed
 105  		}
 106  	}
 107  }
 108  
 109  func (rw *RWCancel) Cancel() (err error) {
 110  	_, err = rw.closingWriter.Write([]byte{0})
 111  	return
 112  }
 113  
 114  func (rw *RWCancel) Close() {
 115  	rw.closingReader.Close()
 116  	rw.closingWriter.Close()
 117  }
 118