85 lines
1.9 KiB
Go
85 lines
1.9 KiB
Go
|
// Copyright (c) HashiCorp, Inc.
|
||
|
|
// SPDX-License-Identifier: MPL-2.0
|
||
|
|
|
||
|
|
package internal
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"net"
|
||
|
|
"sync/atomic"
|
||
|
|
)
|
||
|
|
|
||
|
|
// ErrPipeClosed is returned when calling Accept or DialContext on a closed
|
||
|
|
// PipeListener.
|
||
|
|
var ErrPipeClosed = errors.New("pipe listener has been closed")
|
||
|
|
|
||
|
|
// PipeListener implements the net.Listener interface using a net.Pipe so that
|
||
|
|
// you can interact with a gRPC service in the same process without going over
|
||
|
|
// the network.
|
||
|
|
type PipeListener struct {
|
||
|
|
conns chan net.Conn
|
||
|
|
closed atomic.Bool
|
||
|
|
done chan struct{}
|
||
|
|
}
|
||
|
|
|
||
|
|
var _ net.Listener = (*PipeListener)(nil)
|
||
|
|
|
||
|
|
// NewPipeListener creates a new PipeListener.
|
||
|
|
func NewPipeListener() *PipeListener {
|
||
|
|
return &PipeListener{
|
||
|
|
conns: make(chan net.Conn),
|
||
|
|
done: make(chan struct{}),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Accept a connection.
|
||
|
|
func (p *PipeListener) Accept() (net.Conn, error) {
|
||
|
|
select {
|
||
|
|
case conn := <-p.conns:
|
||
|
|
return conn, nil
|
||
|
|
case <-p.done:
|
||
|
|
return nil, ErrPipeClosed
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Close the listener.
|
||
|
|
func (p *PipeListener) Close() error {
|
||
|
|
if p.closed.CompareAndSwap(false, true) {
|
||
|
|
close(p.done)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// DialContext dials the server over an in-process pipe.
|
||
|
|
func (p *PipeListener) DialContext(ctx context.Context, _ string) (net.Conn, error) {
|
||
|
|
if p.closed.Load() {
|
||
|
|
return nil, ErrPipeClosed
|
||
|
|
}
|
||
|
|
|
||
|
|
serverConn, clientConn := net.Pipe()
|
||
|
|
|
||
|
|
select {
|
||
|
|
// Send the server connection to whatever is accepting connections from the
|
||
|
|
// PipeListener. This will block until something has accepted the conn.
|
||
|
|
case p.conns <- serverConn:
|
||
|
|
return clientConn, nil
|
||
|
|
case <-ctx.Done():
|
||
|
|
serverConn.Close()
|
||
|
|
clientConn.Close()
|
||
|
|
return nil, ctx.Err()
|
||
|
|
case <-p.done:
|
||
|
|
serverConn.Close()
|
||
|
|
clientConn.Close()
|
||
|
|
return nil, ErrPipeClosed
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Add returns the listener's address.
|
||
|
|
func (*PipeListener) Addr() net.Addr { return pipeAddr{} }
|
||
|
|
|
||
|
|
type pipeAddr struct{}
|
||
|
|
|
||
|
|
func (pipeAddr) Network() string { return "pipe" }
|
||
|
|
func (pipeAddr) String() string { return "pipe" }
|