forked from golang/tools
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
internal/fakenet: add a fake network connection
This adds the ability to make a net.Conn based on a reader/writer pair. This is primarily useful for pretending stdin/stdout are a network connection, but can be used for any reader/writer pair. Change-Id: Ib398e62e0fac423db04aed27cfc48070bda2b93b Reviewed-on: https://go-review.googlesource.com/c/tools/+/231697 Run-TryBot: Ian Cottrell <[email protected]> TryBot-Result: Gobot Gobot <[email protected]> Reviewed-by: Robert Findley <[email protected]>
- Loading branch information
Showing
1 changed file
with
129 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,129 @@ | ||
// Copyright 2018 The Go Authors. All rights reserved. | ||
// Use of this source code is governed by a BSD-style | ||
// license that can be found in the LICENSE file. | ||
|
||
package fakenet | ||
|
||
import ( | ||
"io" | ||
"net" | ||
"sync" | ||
"time" | ||
) | ||
|
||
// NewConn returns a net.Conn built on top of the supplied reader and writer. | ||
// It decouples the read and write on the conn from the underlying stream | ||
// to enable Close to abort ones that are in progress. | ||
// It's primary use is to fake a network connection from stdin and stdout. | ||
func NewConn(name string, in io.ReadCloser, out io.WriteCloser) net.Conn { | ||
c := &fakeConn{ | ||
name: name, | ||
reader: newFeeder(in.Read), | ||
writer: newFeeder(out.Write), | ||
in: in, | ||
out: out, | ||
} | ||
go c.reader.run() | ||
go c.writer.run() | ||
return c | ||
} | ||
|
||
type fakeConn struct { | ||
name string | ||
reader *connFeeder | ||
writer *connFeeder | ||
in io.ReadCloser | ||
out io.WriteCloser | ||
} | ||
|
||
type fakeAddr string | ||
|
||
// connFeeder serializes calls to the source function (io.Reader.Read or | ||
// io.Writer.Write) by delegating them to a channel. This also allows calls to | ||
// be intercepted when the connection is closed, and cancelled early if the | ||
// connection is closed while the calls are still outstanding. | ||
type connFeeder struct { | ||
source func([]byte) (int, error) | ||
input chan []byte | ||
result chan feedResult | ||
mu sync.Mutex | ||
closed bool | ||
done chan struct{} | ||
} | ||
|
||
type feedResult struct { | ||
n int | ||
err error | ||
} | ||
|
||
func (c *fakeConn) Close() error { | ||
c.reader.close() | ||
c.writer.close() | ||
c.in.Close() | ||
c.out.Close() | ||
return nil | ||
} | ||
|
||
func (c *fakeConn) Read(b []byte) (n int, err error) { return c.reader.do(b) } | ||
func (c *fakeConn) Write(b []byte) (n int, err error) { return c.writer.do(b) } | ||
func (c *fakeConn) LocalAddr() net.Addr { return fakeAddr(c.name) } | ||
func (c *fakeConn) RemoteAddr() net.Addr { return fakeAddr(c.name) } | ||
func (c *fakeConn) SetDeadline(t time.Time) error { return nil } | ||
func (c *fakeConn) SetReadDeadline(t time.Time) error { return nil } | ||
func (c *fakeConn) SetWriteDeadline(t time.Time) error { return nil } | ||
func (a fakeAddr) Network() string { return "fake" } | ||
func (a fakeAddr) String() string { return string(a) } | ||
|
||
func newFeeder(source func([]byte) (int, error)) *connFeeder { | ||
return &connFeeder{ | ||
source: source, | ||
input: make(chan []byte), | ||
result: make(chan feedResult), | ||
done: make(chan struct{}), | ||
} | ||
} | ||
|
||
func (f *connFeeder) close() { | ||
f.mu.Lock() | ||
if !f.closed { | ||
f.closed = true | ||
close(f.done) | ||
} | ||
f.mu.Unlock() | ||
} | ||
|
||
func (f *connFeeder) do(b []byte) (n int, err error) { | ||
// send the request to the worker | ||
select { | ||
case f.input <- b: | ||
case <-f.done: | ||
return 0, io.EOF | ||
} | ||
// get the result from the worker | ||
select { | ||
case r := <-f.result: | ||
return r.n, r.err | ||
case <-f.done: | ||
return 0, io.EOF | ||
} | ||
} | ||
|
||
func (f *connFeeder) run() { | ||
var b []byte | ||
for { | ||
// wait for an input request | ||
select { | ||
case b = <-f.input: | ||
case <-f.done: | ||
return | ||
} | ||
// invoke the underlying method | ||
n, err := f.source(b) | ||
// send the result back to the requester | ||
select { | ||
case f.result <- feedResult{n: n, err: err}: | ||
case <-f.done: | ||
return | ||
} | ||
} | ||
} |