github.com/amnezia-vpn/amnezia-wg@v0.1.8/rwcancel/rwcancel.go (about)

     1  //go:build !windows && !wasm
     2  
     3  /* SPDX-License-Identifier: MIT
     4   *
     5   * Copyright (C) 2017-2023 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.POLLOUT}}
    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  }