tusd/vendor/google.golang.org/appengine/socket/socket_classic.go

291 lines
7.6 KiB
Go

// Copyright 2012 Google Inc. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.
// +build appengine
package socket
import (
"fmt"
"io"
"net"
"strconv"
"time"
"github.com/golang/protobuf/proto"
"golang.org/x/net/context"
"google.golang.org/appengine/internal"
pb "google.golang.org/appengine/internal/socket"
)
// Dial connects to the address addr on the network protocol.
// The address format is host:port, where host may be a hostname or an IP address.
// Known protocols are "tcp" and "udp".
// The returned connection satisfies net.Conn, and is valid while ctx is valid;
// if the connection is to be used after ctx becomes invalid, invoke SetContext
// with the new context.
func Dial(ctx context.Context, protocol, addr string) (*Conn, error) {
return DialTimeout(ctx, protocol, addr, 0)
}
var ipFamilies = []pb.CreateSocketRequest_SocketFamily{
pb.CreateSocketRequest_IPv4,
pb.CreateSocketRequest_IPv6,
}
// DialTimeout is like Dial but takes a timeout.
// The timeout includes name resolution, if required.
func DialTimeout(ctx context.Context, protocol, addr string, timeout time.Duration) (*Conn, error) {
dialCtx := ctx // Used for dialing and name resolution, but not stored in the *Conn.
if timeout > 0 {
var cancel context.CancelFunc
dialCtx, cancel = context.WithTimeout(ctx, timeout)
defer cancel()
}
host, portStr, err := net.SplitHostPort(addr)
if err != nil {
return nil, err
}
port, err := strconv.Atoi(portStr)
if err != nil {
return nil, fmt.Errorf("socket: bad port %q: %v", portStr, err)
}
var prot pb.CreateSocketRequest_SocketProtocol
switch protocol {
case "tcp":
prot = pb.CreateSocketRequest_TCP
case "udp":
prot = pb.CreateSocketRequest_UDP
default:
return nil, fmt.Errorf("socket: unknown protocol %q", protocol)
}
packedAddrs, resolved, err := resolve(dialCtx, ipFamilies, host)
if err != nil {
return nil, fmt.Errorf("socket: failed resolving %q: %v", host, err)
}
if len(packedAddrs) == 0 {
return nil, fmt.Errorf("no addresses for %q", host)
}
packedAddr := packedAddrs[0] // use first address
fam := pb.CreateSocketRequest_IPv4
if len(packedAddr) == net.IPv6len {
fam = pb.CreateSocketRequest_IPv6
}
req := &pb.CreateSocketRequest{
Family: fam.Enum(),
Protocol: prot.Enum(),
RemoteIp: &pb.AddressPort{
Port: proto.Int32(int32(port)),
PackedAddress: packedAddr,
},
}
if resolved {
req.RemoteIp.HostnameHint = &host
}
res := &pb.CreateSocketReply{}
if err := internal.Call(dialCtx, "remote_socket", "CreateSocket", req, res); err != nil {
return nil, err
}
return &Conn{
ctx: ctx,
desc: res.GetSocketDescriptor(),
prot: prot,
local: res.ProxyExternalIp,
remote: req.RemoteIp,
}, nil
}
// LookupIP returns the given host's IP addresses.
func LookupIP(ctx context.Context, host string) (addrs []net.IP, err error) {
packedAddrs, _, err := resolve(ctx, ipFamilies, host)
if err != nil {
return nil, fmt.Errorf("socket: failed resolving %q: %v", host, err)
}
addrs = make([]net.IP, len(packedAddrs))
for i, pa := range packedAddrs {
addrs[i] = net.IP(pa)
}
return addrs, nil
}
func resolve(ctx context.Context, fams []pb.CreateSocketRequest_SocketFamily, host string) ([][]byte, bool, error) {
// Check if it's an IP address.
if ip := net.ParseIP(host); ip != nil {
if ip := ip.To4(); ip != nil {
return [][]byte{ip}, false, nil
}
return [][]byte{ip}, false, nil
}
req := &pb.ResolveRequest{
Name: &host,
AddressFamilies: fams,
}
res := &pb.ResolveReply{}
if err := internal.Call(ctx, "remote_socket", "Resolve", req, res); err != nil {
// XXX: need to map to pb.ResolveReply_ErrorCode?
return nil, false, err
}
return res.PackedAddress, true, nil
}
// withDeadline is like context.WithDeadline, except it ignores the zero deadline.
func withDeadline(parent context.Context, deadline time.Time) (context.Context, context.CancelFunc) {
if deadline.IsZero() {
return parent, func() {}
}
return context.WithDeadline(parent, deadline)
}
// Conn represents a socket connection.
// It implements net.Conn.
type Conn struct {
ctx context.Context
desc string
offset int64
prot pb.CreateSocketRequest_SocketProtocol
local, remote *pb.AddressPort
readDeadline, writeDeadline time.Time // optional
}
// SetContext sets the context that is used by this Conn.
// It is usually used only when using a Conn that was created in a different context,
// such as when a connection is created during a warmup request but used while
// servicing a user request.
func (cn *Conn) SetContext(ctx context.Context) {
cn.ctx = ctx
}
func (cn *Conn) Read(b []byte) (n int, err error) {
const maxRead = 1 << 20
if len(b) > maxRead {
b = b[:maxRead]
}
req := &pb.ReceiveRequest{
SocketDescriptor: &cn.desc,
DataSize: proto.Int32(int32(len(b))),
}
res := &pb.ReceiveReply{}
if !cn.readDeadline.IsZero() {
req.TimeoutSeconds = proto.Float64(cn.readDeadline.Sub(time.Now()).Seconds())
}
ctx, cancel := withDeadline(cn.ctx, cn.readDeadline)
defer cancel()
if err := internal.Call(ctx, "remote_socket", "Receive", req, res); err != nil {
return 0, err
}
if len(res.Data) == 0 {
return 0, io.EOF
}
if len(res.Data) > len(b) {
return 0, fmt.Errorf("socket: internal error: read too much data: %d > %d", len(res.Data), len(b))
}
return copy(b, res.Data), nil
}
func (cn *Conn) Write(b []byte) (n int, err error) {
const lim = 1 << 20 // max per chunk
for n < len(b) {
chunk := b[n:]
if len(chunk) > lim {
chunk = chunk[:lim]
}
req := &pb.SendRequest{
SocketDescriptor: &cn.desc,
Data: chunk,
StreamOffset: &cn.offset,
}
res := &pb.SendReply{}
if !cn.writeDeadline.IsZero() {
req.TimeoutSeconds = proto.Float64(cn.writeDeadline.Sub(time.Now()).Seconds())
}
ctx, cancel := withDeadline(cn.ctx, cn.writeDeadline)
defer cancel()
if err = internal.Call(ctx, "remote_socket", "Send", req, res); err != nil {
// assume zero bytes were sent in this RPC
break
}
n += int(res.GetDataSent())
cn.offset += int64(res.GetDataSent())
}
return
}
func (cn *Conn) Close() error {
req := &pb.CloseRequest{
SocketDescriptor: &cn.desc,
}
res := &pb.CloseReply{}
if err := internal.Call(cn.ctx, "remote_socket", "Close", req, res); err != nil {
return err
}
cn.desc = "CLOSED"
return nil
}
func addr(prot pb.CreateSocketRequest_SocketProtocol, ap *pb.AddressPort) net.Addr {
if ap == nil {
return nil
}
switch prot {
case pb.CreateSocketRequest_TCP:
return &net.TCPAddr{
IP: net.IP(ap.PackedAddress),
Port: int(*ap.Port),
}
case pb.CreateSocketRequest_UDP:
return &net.UDPAddr{
IP: net.IP(ap.PackedAddress),
Port: int(*ap.Port),
}
}
panic("unknown protocol " + prot.String())
}
func (cn *Conn) LocalAddr() net.Addr { return addr(cn.prot, cn.local) }
func (cn *Conn) RemoteAddr() net.Addr { return addr(cn.prot, cn.remote) }
func (cn *Conn) SetDeadline(t time.Time) error {
cn.readDeadline = t
cn.writeDeadline = t
return nil
}
func (cn *Conn) SetReadDeadline(t time.Time) error {
cn.readDeadline = t
return nil
}
func (cn *Conn) SetWriteDeadline(t time.Time) error {
cn.writeDeadline = t
return nil
}
// KeepAlive signals that the connection is still in use.
// It may be called to prevent the socket being closed due to inactivity.
func (cn *Conn) KeepAlive() error {
req := &pb.GetSocketNameRequest{
SocketDescriptor: &cn.desc,
}
res := &pb.GetSocketNameReply{}
return internal.Call(cn.ctx, "remote_socket", "GetSocketName", req, res)
}
func init() {
internal.RegisterErrorCodeMap("remote_socket", pb.RemoteSocketServiceError_ErrorCode_name)
}