github.com/oarkflow/log@v1.0.78/console_windows.go (about)

     1  //go:build windows
     2  // +build windows
     3  
     4  package log
     5  
     6  import (
     7  	"io"
     8  	"os"
     9  	"sync"
    10  	"syscall"
    11  	"unsafe"
    12  )
    13  
    14  func isTerminal(fd uintptr, _, _ string) bool {
    15  	var mode uint32
    16  	err := syscall.GetConsoleMode(syscall.Handle(fd), &mode)
    17  	if err != nil {
    18  		return false
    19  	}
    20  
    21  	return true
    22  }
    23  
    24  var (
    25  	kernel32                = syscall.NewLazyDLL("kernel32.dll")
    26  	setConsoleMode          = kernel32.NewProc("SetConsoleMode").Call
    27  	setConsoleTextAttribute = kernel32.NewProc("SetConsoleTextAttribute").Call
    28  
    29  	muConsole   sync.Mutex
    30  	onceConsole sync.Once
    31  	isvt        bool
    32  )
    33  
    34  // WriteEntry implements Writer
    35  func (w *ConsoleWriter) WriteEntry(e *Entry) (n int, err error) {
    36  	onceConsole.Do(func() { isvt = isVirtualTerminal() })
    37  
    38  	out := w.Writer
    39  	if out == nil {
    40  		out = os.Stderr
    41  	}
    42  	if isvt {
    43  		n, err = w.write(out, e.buf)
    44  	} else {
    45  		n, err = w.writew(out, e.buf)
    46  	}
    47  	return
    48  }
    49  
    50  func (w *ConsoleWriter) writew(out io.Writer, p []byte) (n int, err error) {
    51  	muConsole.Lock()
    52  	defer muConsole.Unlock()
    53  
    54  	b := bbpool.Get().(*bb)
    55  	b.B = b.B[:0]
    56  	defer bbpool.Put(b)
    57  
    58  	n, err = w.write(b, p)
    59  	if err != nil {
    60  		return
    61  	}
    62  	n = 0
    63  	// uintptr color
    64  	const (
    65  		Black  = 0
    66  		Blue   = 1
    67  		Green  = 2
    68  		Aqua   = 3
    69  		Red    = 4
    70  		Purple = 5
    71  		Yellow = 6
    72  		White  = 7
    73  		Gray   = 8
    74  	)
    75  	// color print
    76  	var cprint = func(color uintptr, b []byte) {
    77  		if color != White {
    78  			setConsoleTextAttribute(uintptr(syscall.Stderr), color)
    79  			defer setConsoleTextAttribute(uintptr(syscall.Stderr), White)
    80  		}
    81  		var i int
    82  		i, err = out.Write(b)
    83  		n += i
    84  	}
    85  
    86  	b2 := bbpool.Get().(*bb)
    87  	b2.B = b2.B[:0]
    88  	defer bbpool.Put(b2)
    89  
    90  	var color uintptr = White
    91  	var length = len(b.B)
    92  	var c uint32
    93  	for i := 0; i < length; i++ {
    94  		if b.B[i] == '\x1b' {
    95  			switch {
    96  			case length-i > 3 &&
    97  				b.B[i+1] == '[' &&
    98  				'0' <= b.B[i+2] && b.B[i+2] <= '9' &&
    99  				b.B[i+3] == 'm':
   100  				c = uint32(b.B[i+2] - '0')
   101  				i += 3
   102  			case length-i > 4 &&
   103  				b.B[i+1] == '[' &&
   104  				'0' <= b.B[i+2] && b.B[i+2] <= '9' &&
   105  				'0' <= b.B[i+3] && b.B[i+3] <= '9' &&
   106  				b.B[i+4] == 'm':
   107  				c = uint32(b.B[i+2]-'0')*10 + uint32(b.B[i+3]-'0')
   108  				i += 4
   109  			}
   110  			if len(b2.B) > 0 {
   111  				cprint(color, b2.B)
   112  			}
   113  			b2.B = b2.B[:0]
   114  			switch c {
   115  			case 0: // Reset
   116  				color = White
   117  			case 30: // Black
   118  				color = Black
   119  			case 90: // Gray
   120  				color = Gray
   121  			case 31, 91: // Red, BrightRed
   122  				color = Red
   123  			case 32, 92: // Green, BrightGreen
   124  				color = Green
   125  			case 33, 93: // Yellow, BrightYellow
   126  				color = Yellow
   127  			case 34, 94: // Blue, BrightBlue
   128  				color = Blue
   129  			case 35, 95: // Magenta, BrightMagenta
   130  				color = Purple
   131  			case 36, 96: // Cyan, BrightCyan
   132  				color = Aqua
   133  			case 37, 97: // White, BrightWhite
   134  				color = White
   135  			default:
   136  				color = White
   137  			}
   138  		} else {
   139  			b2.B = append(b2.B, b.B[i])
   140  		}
   141  	}
   142  
   143  	if len(b2.B) != 0 {
   144  		cprint(White, b2.B)
   145  	}
   146  
   147  	return
   148  }
   149  
   150  func isVirtualTerminal() bool {
   151  	var h syscall.Handle
   152  	var b [64]uint16
   153  	var n uint32
   154  
   155  	// open registry
   156  	err := syscall.RegOpenKeyEx(syscall.HKEY_LOCAL_MACHINE, syscall.StringToUTF16Ptr(`SOFTWARE\Microsoft\Windows NT\CurrentVersion`), 0, syscall.KEY_READ, &h)
   157  	if err != nil {
   158  		return false
   159  	}
   160  	defer syscall.RegCloseKey(h)
   161  
   162  	// read windows build number
   163  	n = uint32(len(b))
   164  	err = syscall.RegQueryValueEx(h, syscall.StringToUTF16Ptr(`CurrentBuild`), nil, nil, (*byte)(unsafe.Pointer(&b[0])), &n)
   165  	if err != nil {
   166  		return false
   167  	}
   168  	for i := 0; i < len(b); i++ {
   169  		if b[i] == 0 {
   170  			break
   171  		}
   172  		n = n*10 + uint32(b[i]-'0')
   173  	}
   174  
   175  	// return if lower than windows 10 16257
   176  	if n < 16257 {
   177  		return false
   178  	}
   179  
   180  	// get console mode
   181  	err = syscall.GetConsoleMode(syscall.Stderr, &n)
   182  	if err != nil {
   183  		return false
   184  	}
   185  
   186  	// enable ENABLE_VIRTUAL_TERMINAL_PROCESSING
   187  	ret, _, _ := setConsoleMode(uintptr(syscall.Stderr), uintptr(n|0x4))
   188  	return ret != 0
   189  }