github.com/moontrade/mdbx-go@v0.4.0/internal/capture/capture.go (about)

     1  package capture
     2  
     3  /*
     4  #include <stdio.h>
     5  #include <stdlib.h>
     6  */
     7  import "C"
     8  
     9  import (
    10  	"bufio"
    11  	"bytes"
    12  	"io"
    13  	"os"
    14  	"sync"
    15  	"syscall"
    16  )
    17  
    18  var lockStdFileDescriptorsSwapping sync.Mutex
    19  
    20  // Capture captures stderr and stdout of a given function call.
    21  func Capture(call func()) (output []byte, err error) {
    22  	originalStdErr, originalStdOut := os.Stderr, os.Stdout
    23  	defer func() {
    24  		lockStdFileDescriptorsSwapping.Lock()
    25  
    26  		os.Stderr, os.Stdout = originalStdErr, originalStdOut
    27  
    28  		lockStdFileDescriptorsSwapping.Unlock()
    29  	}()
    30  
    31  	r, w, err := os.Pipe()
    32  	if err != nil {
    33  		return nil, err
    34  	}
    35  	defer func() {
    36  		e := r.Close()
    37  		if e != nil {
    38  			err = e
    39  		}
    40  		if w != nil {
    41  			e = w.Close()
    42  			if err != nil {
    43  				err = e
    44  			}
    45  		}
    46  	}()
    47  
    48  	lockStdFileDescriptorsSwapping.Lock()
    49  
    50  	os.Stderr, os.Stdout = w, w
    51  
    52  	lockStdFileDescriptorsSwapping.Unlock()
    53  
    54  	out := make(chan []byte)
    55  	go func() {
    56  		defer func() {
    57  			// If there is a panic in the function call, copying from "r" does not work anymore.
    58  			_ = recover()
    59  		}()
    60  
    61  		var b bytes.Buffer
    62  
    63  		_, err := io.Copy(&b, r)
    64  		if err != nil {
    65  			panic(err)
    66  		}
    67  
    68  		out <- b.Bytes()
    69  	}()
    70  
    71  	call()
    72  
    73  	err = w.Close()
    74  	if err != nil {
    75  		return nil, err
    76  	}
    77  	w = nil
    78  
    79  	return <-out, err
    80  }
    81  
    82  // CaptureWithCGo captures stderr and stdout as well as stderr and stdout of C of a given function call.
    83  func CaptureWithCGo(call func()) (output []byte, err error) {
    84  	lockStdFileDescriptorsSwapping.Lock()
    85  
    86  	originalStdout, e := syscall.Dup(syscall.Stdout)
    87  	if e != nil {
    88  		lockStdFileDescriptorsSwapping.Unlock()
    89  
    90  		return nil, e
    91  	}
    92  
    93  	originalStderr, e := syscall.Dup(syscall.Stderr)
    94  	if e != nil {
    95  		lockStdFileDescriptorsSwapping.Unlock()
    96  
    97  		return nil, e
    98  	}
    99  
   100  	lockStdFileDescriptorsSwapping.Unlock()
   101  
   102  	defer func() {
   103  		lockStdFileDescriptorsSwapping.Lock()
   104  
   105  		if e := syscall.Dup2(originalStdout, syscall.Stdout); e != nil {
   106  			lockStdFileDescriptorsSwapping.Unlock()
   107  
   108  			err = e
   109  		}
   110  		if e := syscall.Close(originalStdout); e != nil {
   111  			lockStdFileDescriptorsSwapping.Unlock()
   112  
   113  			err = e
   114  		}
   115  		if e := syscall.Dup2(originalStderr, syscall.Stderr); e != nil {
   116  			lockStdFileDescriptorsSwapping.Unlock()
   117  
   118  			err = e
   119  		}
   120  		if e := syscall.Close(originalStderr); e != nil {
   121  			lockStdFileDescriptorsSwapping.Unlock()
   122  
   123  			err = e
   124  		}
   125  
   126  		lockStdFileDescriptorsSwapping.Unlock()
   127  	}()
   128  
   129  	r, w, err := os.Pipe()
   130  	if err != nil {
   131  		return nil, err
   132  	}
   133  	defer func() {
   134  		e := r.Close()
   135  		if e != nil {
   136  			err = e
   137  		}
   138  		if w != nil {
   139  			e = w.Close()
   140  			if err != nil {
   141  				err = e
   142  			}
   143  		}
   144  	}()
   145  
   146  	lockStdFileDescriptorsSwapping.Lock()
   147  
   148  	if e := syscall.Dup2(int(w.Fd()), syscall.Stdout); e != nil {
   149  		lockStdFileDescriptorsSwapping.Unlock()
   150  
   151  		return nil, e
   152  	}
   153  	if e := syscall.Dup2(int(w.Fd()), syscall.Stderr); e != nil {
   154  		lockStdFileDescriptorsSwapping.Unlock()
   155  
   156  		return nil, e
   157  	}
   158  
   159  	lockStdFileDescriptorsSwapping.Unlock()
   160  
   161  	out := make(chan []byte)
   162  	go func() {
   163  		defer func() {
   164  			// If there is a panic in the function call, copying from "r" does not work anymore.
   165  			_ = recover()
   166  		}()
   167  
   168  		var b bytes.Buffer
   169  
   170  		_, err := io.Copy(&b, r)
   171  		if err != nil {
   172  			panic(err)
   173  		}
   174  
   175  		out <- b.Bytes()
   176  	}()
   177  
   178  	call()
   179  
   180  	lockStdFileDescriptorsSwapping.Lock()
   181  
   182  	C.fflush(C.stdout)
   183  
   184  	err = w.Close()
   185  	if err != nil {
   186  		lockStdFileDescriptorsSwapping.Unlock()
   187  
   188  		return nil, err
   189  	}
   190  	w = nil
   191  
   192  	if e := syscall.Close(syscall.Stdout); e != nil {
   193  		lockStdFileDescriptorsSwapping.Unlock()
   194  
   195  		return nil, e
   196  	}
   197  	if e := syscall.Close(syscall.Stderr); e != nil {
   198  		lockStdFileDescriptorsSwapping.Unlock()
   199  
   200  		return nil, e
   201  	}
   202  
   203  	lockStdFileDescriptorsSwapping.Unlock()
   204  
   205  	return <-out, err
   206  }
   207  
   208  type Msg struct {
   209  	Line string
   210  }
   211  
   212  // CaptureWithCGo captures stderr and stdout as well as stderr and stdout of C of a given function call.
   213  func CaptureWithCGoChan(ch chan string, call func()) (err error) {
   214  	lockStdFileDescriptorsSwapping.Lock()
   215  
   216  	originalStdout, e := syscall.Dup(syscall.Stdout)
   217  	if e != nil {
   218  		lockStdFileDescriptorsSwapping.Unlock()
   219  
   220  		return e
   221  	}
   222  
   223  	originalStderr, e := syscall.Dup(syscall.Stderr)
   224  	if e != nil {
   225  		lockStdFileDescriptorsSwapping.Unlock()
   226  
   227  		return e
   228  	}
   229  
   230  	lockStdFileDescriptorsSwapping.Unlock()
   231  
   232  	defer func() {
   233  		lockStdFileDescriptorsSwapping.Lock()
   234  
   235  		if e := syscall.Dup2(originalStdout, syscall.Stdout); e != nil {
   236  			lockStdFileDescriptorsSwapping.Unlock()
   237  
   238  			err = e
   239  		}
   240  		if e := syscall.Close(originalStdout); e != nil {
   241  			lockStdFileDescriptorsSwapping.Unlock()
   242  
   243  			err = e
   244  		}
   245  		if e := syscall.Dup2(originalStderr, syscall.Stderr); e != nil {
   246  			lockStdFileDescriptorsSwapping.Unlock()
   247  
   248  			err = e
   249  		}
   250  		if e := syscall.Close(originalStderr); e != nil {
   251  			lockStdFileDescriptorsSwapping.Unlock()
   252  
   253  			err = e
   254  		}
   255  
   256  		lockStdFileDescriptorsSwapping.Unlock()
   257  	}()
   258  
   259  	r, w, err := os.Pipe()
   260  	if err != nil {
   261  		return err
   262  	}
   263  	defer func() {
   264  		e := r.Close()
   265  		if e != nil {
   266  			err = e
   267  		}
   268  		if w != nil {
   269  			e = w.Close()
   270  			if err != nil {
   271  				err = e
   272  			}
   273  		}
   274  	}()
   275  
   276  	lockStdFileDescriptorsSwapping.Lock()
   277  
   278  	if e := syscall.Dup2(int(w.Fd()), syscall.Stdout); e != nil {
   279  		lockStdFileDescriptorsSwapping.Unlock()
   280  
   281  		return e
   282  	}
   283  	if e := syscall.Dup2(int(w.Fd()), syscall.Stderr); e != nil {
   284  		lockStdFileDescriptorsSwapping.Unlock()
   285  
   286  		return e
   287  	}
   288  
   289  	lockStdFileDescriptorsSwapping.Unlock()
   290  
   291  	wg := &sync.WaitGroup{}
   292  	wg.Add(1)
   293  	go func() {
   294  		defer func() {
   295  			wg.Done()
   296  			// If there is a panic in the function call, copying from "r" does not work anymore.
   297  			_ = recover()
   298  		}()
   299  		scanner := bufio.NewScanner(r)
   300  		for scanner.Scan() {
   301  			text := scanner.Text()
   302  			ch <- text
   303  		}
   304  	}()
   305  
   306  	call()
   307  
   308  	lockStdFileDescriptorsSwapping.Lock()
   309  
   310  	C.fflush(C.stdout)
   311  
   312  	err = w.Close()
   313  	if err != nil {
   314  		lockStdFileDescriptorsSwapping.Unlock()
   315  
   316  		return err
   317  	}
   318  	w = nil
   319  
   320  	if e := syscall.Close(syscall.Stdout); e != nil {
   321  		lockStdFileDescriptorsSwapping.Unlock()
   322  
   323  		return e
   324  	}
   325  	if e := syscall.Close(syscall.Stderr); e != nil {
   326  		lockStdFileDescriptorsSwapping.Unlock()
   327  
   328  		return e
   329  	}
   330  
   331  	lockStdFileDescriptorsSwapping.Unlock()
   332  
   333  	wg.Wait()
   334  	return err
   335  }
   336  
   337  type Buffer struct {
   338  	rd *io.PipeReader
   339  	wr *io.PipeWriter
   340  }
   341  
   342  func NewBuffer() *Buffer {
   343  	rd, wr := io.Pipe()
   344  	return &Buffer{rd, wr}
   345  }
   346  
   347  func (b *Buffer) Close() error {
   348  	if b.rd == nil {
   349  		return nil
   350  	}
   351  	_ = b.rd.Close()
   352  	_ = b.wr.Close()
   353  	b.rd = nil
   354  	b.wr = nil
   355  	return nil
   356  }
   357  
   358  func (b *Buffer) Read(d []byte) (int, error) {
   359  	return b.rd.Read(d)
   360  }
   361  
   362  func (b *Buffer) Write(d []byte) (int, error) {
   363  	return b.wr.Write(d)
   364  }