github.com/MetalBlockchain/metalgo@v1.11.9/vms/rpcchainvm/ghttp/gresponsewriter/locked_writer.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package gresponsewriter
     5  
     6  import (
     7  	"bufio"
     8  	"net"
     9  	"net/http"
    10  	"sync"
    11  )
    12  
    13  var (
    14  	_ http.ResponseWriter = (*lockedWriter)(nil)
    15  	_ http.Flusher        = (*lockedWriter)(nil)
    16  	_ http.Hijacker       = (*lockedWriter)(nil)
    17  )
    18  
    19  type lockedWriter struct {
    20  	lock          sync.Mutex
    21  	writer        http.ResponseWriter
    22  	headerWritten bool
    23  }
    24  
    25  func NewLockedWriter(w http.ResponseWriter) http.ResponseWriter {
    26  	return &lockedWriter{writer: w}
    27  }
    28  
    29  func (lw *lockedWriter) Header() http.Header {
    30  	lw.lock.Lock()
    31  	defer lw.lock.Unlock()
    32  
    33  	return lw.writer.Header()
    34  }
    35  
    36  func (lw *lockedWriter) Write(b []byte) (int, error) {
    37  	lw.lock.Lock()
    38  	defer lw.lock.Unlock()
    39  
    40  	lw.headerWritten = true
    41  	return lw.writer.Write(b)
    42  }
    43  
    44  func (lw *lockedWriter) WriteHeader(statusCode int) {
    45  	lw.lock.Lock()
    46  	defer lw.lock.Unlock()
    47  
    48  	// Skip writing the header if it has already been written once.
    49  	if lw.headerWritten {
    50  		return
    51  	}
    52  	lw.headerWritten = true
    53  	lw.writer.WriteHeader(statusCode)
    54  }
    55  
    56  func (lw *lockedWriter) Flush() {
    57  	lw.lock.Lock()
    58  	defer lw.lock.Unlock()
    59  
    60  	flusher, ok := lw.writer.(http.Flusher)
    61  	if ok {
    62  		flusher.Flush()
    63  	}
    64  }
    65  
    66  func (lw *lockedWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    67  	lw.lock.Lock()
    68  	defer lw.lock.Unlock()
    69  
    70  	hijacker, ok := lw.writer.(http.Hijacker)
    71  	if !ok {
    72  		return nil, nil, errUnsupportedHijacking
    73  	}
    74  	return hijacker.Hijack()
    75  }