<commit_before>/*
Copyright 2016 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package tensorflow

// #include "tensorflow/c/c_api.h"
//
// #include <stdlib.h>
// #include <string.h>
//
// void TF_SetAttrShapeList_Helper(TF_OperationDescription* desc,
//                                 const char* attr_name,
//                                 const int64_t* flat_dims,
//                                 const int* num_dims,
//                                 int num_shapes) {
//  const int64_t** dims =
//    (const int64_t**)malloc(sizeof(const int64_t*) * num_shapes);
//  int i = 0;
//  for (i = 0; i < num_shapes; i++) {
//    dims[i] = flat_dims;
//    if (num_dims[i] > 0) {
//      // flat_dims will be NULL iff num_shapes is 0 or all elements in num_dims are <= 0.
//      flat_dims += num_dims[i];
//    }
//  }
//  TF_SetAttrShapeList(desc, attr_name, dims, num_dims, num_shapes);
//  free(dims);
// }
import "C"

import (
	"fmt"
	"io"
	"runtime"
	"unsafe"
)

// Graph represents a computation graph. Graphs may be shared between sessions.
type Graph struct {
	c *C.TF_Graph
}

// NewGraph returns a new Graph.
func NewGraph() *Graph {
	g := &Graph{C.TF_NewGraph()}
	runtime.SetFinalizer(g, (*Graph).finalizer)
	return g
}

func (g *Graph) finalizer() {
	C.TF_DeleteGraph(g.c)
}

// WriteTo writes out a serialized representation of g to w.
//
// Implements the io.WriterTo interface.
func (g *Graph) WriteTo(w io.Writer) (int64, error) {
	buf := C.TF_NewBuffer()
	defer C.TF_DeleteBuffer(buf)
	status := newStatus()
	C.TF_GraphToGraphDef(g.c, buf, status.c)
	if err := status.Err(); err != nil {
		return 0, err
	}
	if buf.length > (1 << 30) {
		// For very large graphs, the writes can be chunked.
		// Punt on that for now.
		return 0, fmt.Errorf("Graph is too large to write out, Graph.WriteTo needs to be updated")
	}
	// A []byte slice backed by C memory.
	// See: https://github.com/golang/go/wiki/cgo#turning-c-arrays-into-go-slices
	length := int(buf.length)
	slice := (*[1 << 30]byte)(unsafe.Pointer(buf.data))[:length:length]
	n, err := w.Write(slice)
	return int64(n), err
}

// Import imports the nodes and edges from a serialized representation of
// another Graph into g.
//
// Names of imported nodes will be prefixed with prefix.
func (g *Graph) ImportWithDevice(def []byte, prefix string, device string) error {
	cprefix := C.CString(prefix)
	defer C.free(unsafe.Pointer(cprefix))

	opts := C.TF_NewImportGraphDefOptions()
	defer C.TF_DeleteImportGraphDefOptions(opts)
	C.TF_ImportGraphDefOptionsSetPrefix(opts, cprefix)

	if len(device) != 0 {
		cdev := C.CString(device)
		defer C.free(unsafe.Pointer(cdev))
		C.TF_ImportGraphDefOptionsSetBindDevice(opts, cdev)
	}

	buf := C.TF_NewBuffer()
	defer C.TF_DeleteBuffer(buf)
	// Would have preferred to use C.CBytes, but that does not play well
	// with "go vet" till https://github.com/golang/go/issues/17201 is
	// resolved.
	buf.length = C.size_t(len(def))
	buf.data = C.malloc(buf.length)
	if buf.data == nil {
		return fmt.Errorf("unable to allocate memory")
	}
	defer C.free(buf.data)
	C.memcpy(buf.data, unsafe.Pointer(&def[0]), buf.length)

	status := newStatus()

	C.TF_GraphImportGraphDef(g.c, buf, opts, status.c)
	if err := status.Err(); err != nil {
		return err
	}

	return nil
}

func (g *Graph) Import(def []byte, prefix string) error {
    return g.ImportWithDevice(def, prefix, "")
}

// Operation returns the Operation named name in the Graph, or nil if no such
// operation is present.
func (g *Graph) Operation(name string) *Operation {
	cname := C.CString(name)
	defer C.free(unsafe.Pointer(cname))
	cop := C.TF_GraphOperationByName(g.c, cname)
	if cop == nil {
		return nil
	}
	return &Operation{cop, g}
}

// Operations returns a list of all operations in the graph
func (g *Graph) Operations() []Operation {
	var pos C.size_t = 0
	ops := []Operation{}
	for {
		cop := C.TF_GraphNextOperation(g.c, &pos)
		if cop == nil {
			break
		}
		ops = append(ops, Operation{cop, g})
	}
	return ops
}

// OpSpec is the specification of an Operation to be added to a Graph
// (using Graph.AddOperation).
type OpSpec struct {
	// Type of the operation (e.g., "Add", "MatMul").
	Type string

	// Name by which the added operation will be referred to in the Graph.
	// If omitted, defaults to Type.
	Name string

	// Inputs to this operation, which in turn must be outputs
	// of other operations already added to the Graph.
	//
	// An operation may have multiple inputs with individual inputs being
	// either a single tensor produced by another operation or a list of
	// tensors produced by multiple operations. For example, the "Concat"
	// operation takes two inputs: (1) the dimension along which to
	// concatenate and (2) a list of tensors to concatenate. Thus, for
	// Concat, len(Input) must be 2, with the first element being an Output
	// and the second being an OutputList.
	Input []Input

	// Map from attribute name to its value that will be attached to this
	// operation.
	Attrs map[string]interface{}

	// Operations that must be executed before executing the operation
	// being added.
	ControlDependencies []*Operation

	// The device on which the operation should be executed.
	// If omitted, an appropriate device will automatically be selected.
	//
	// For example, if set of "/device:GPU:0", then the operation will
	// execute on GPU #0.
	Device string

	// Other possible fields: ColocateWith.
}

// AddOperation adds an operation to g.
func (g *Graph) AddOperation(args OpSpec) (*Operation, error) {
	if args.Name == "" {
		args.Name = args.Type
	}
	cname := C.CString(args.Name)
	ctype := C.CString(args.Type)
	cdesc := C.TF_NewOperation(g.c, ctype, cname)
	C.free(unsafe.Pointer(cname))
	C.free(unsafe.Pointer(ctype))

	for _, in := range args.Input {
		switch in := in.(type) {
		case Output:
			C.TF_AddInput(cdesc, in.c())
		case OutputList:
			size := len(in)
			list := make([]C.TF_Output, size)
			for i, v := range in {
				list[i] = v.c()
			}
			if size > 0 {
				C.TF_AddInputList(cdesc, &list[0], C.int(size))
			} else {
				C.TF_AddInputList(cdesc, nil, 0)
			}
		}
	}
	for _, in := range args.ControlDependencies {
		C.TF_AddControlInput(cdesc, in.c)
	}
	status := newStatus()
	for name, value := range args.Attrs {
		if err := setAttr(cdesc, status, name, value); err != nil {
			// Memory leak here as the TF_OperationDescription
			// object will not be cleaned up. At the time of this
			// writing, this was next to impossible since it
			// required value to be a string tensor with
			// incorrectly encoded strings. Given this rarity, live
			// with the memory leak.  If it becomes a real problem,
			// consider adding a TF_DeleteOperationDescription
			// function to the C API.
			return nil, fmt.Errorf("%v (memory will be leaked)", err)
		}
	}
	if len(args.Device) > 0 {
		cdevice := C.CString(args.Device)
		C.TF_SetDevice(cdesc, cdevice)
		C.free(unsafe.Pointer(cdevice))
	}
	c := C.TF_FinishOperation(cdesc, status.c)
	if err := status.Err(); err != nil {
		return nil, err
	}
	return &Operation{c, g}, nil
}

func setAttr(cdesc *C.TF_OperationDescription, status *status, name string, value interface{}) error {
	cAttrName := C.CString(name)
	defer C.free(unsafe.Pointer(cAttrName))
	switch value := value.(type) {
	case string:
		cstr := C.CString(value)
		C.TF_SetAttrString(cdesc, cAttrName, unsafe.Pointer(cstr), C.size_t(len(value)))
		C.free(unsafe.Pointer(cstr))
	case []string:
		size := len(value)
		list := make([]unsafe.Pointer, size)
		lens := make([]C.size_t, size)
		for i, s := range value {
			list[i] = unsafe.Pointer(C.CString(s))
			lens[i] = C.size_t(len(s))
		}
		if size > 0 {
			C.TF_SetAttrStringList(cdesc, cAttrName, &list[0], &lens[0], C.int(size))
		} else {
			C.TF_SetAttrStringList(cdesc, cAttrName, nil, nil, 0)
		}
		for _, s := range list {
			C.free(s)
		}
	case int64:
		C.TF_SetAttrInt(cdesc, cAttrName, C.int64_t(value))
	case []int64:
		size := len(value)
		list := make([]C.int64_t, size)
		for i, v := range value {
			list[i] = C.int64_t(v)
		}
		if size > 0 {
			C.TF_SetAttrIntList(cdesc, cAttrName, &list[0], C.int(size))
		} else {
			C.TF_SetAttrIntList(cdesc, cAttrName, nil, 0)
		}
	case float32:
		C.TF_SetAttrFloat(cdesc, cAttrName, C.float(value))
	case []float32:
		size := len(value)
		list := make([]C.float, size)
		for i, v := range value {
			list[i] = C.float(v)
		}
		if size > 0 {
			C.TF_SetAttrFloatList(cdesc, cAttrName, &list[0], C.int(size))
		} else {
			C.TF_SetAttrFloatList(cdesc, cAttrName, nil, 0)
		}
	case bool:
		v := C.uchar(0)
		if value {
			v = 1
		}
		C.TF_SetAttrBool(cdesc, cAttrName, v)
	case []bool:
		size := len(value)
		list := make([]C.uchar, size)
		for i, v := range value {
			if v {
				list[i] = 1
			}
		}
		if size > 0 {
			C.TF_SetAttrBoolList(cdesc, cAttrName, &list[0], C.int(size))
		} else {
			C.TF_SetAttrBoolList(cdesc, cAttrName, nil, 0)
		}
	case DataType:
		C.TF_SetAttrType(cdesc, cAttrName, C.TF_DataType(value))
	case []DataType:
		var list *C.TF_DataType
		if len(value) > 0 {
			list = (*C.TF_DataType)(&value[0])
		}
		C.TF_SetAttrTypeList(cdesc, cAttrName, list, C.int(len(value)))
	case *Tensor:
		C.TF_SetAttrTensor(cdesc, cAttrName, value.c, status.c)
		if err := status.Err(); err != nil {
			return fmt.Errorf("bad value for attribute %q: %v", name, err)
		}
	case []*Tensor:
		size := len(value)
		list := make([]*C.TF_Tensor, size)
		for i, v := range value {
			list[i] = v.c
		}
		var plist **C.TF_Tensor
		if size > 0 {
			plist = &list[0]
		}
		C.TF_SetAttrTensorList(cdesc, cAttrName, plist, C.int(size), status.c)
		if err := status.Err(); err != nil {
			return fmt.Errorf("bad value for attribute %q: %v", name, err)
		}
	case Shape:
		ndims := C.int(value.NumDimensions())
		var dimsp *C.int64_t
		if ndims > 0 {
			dims := make([]C.int64_t, ndims)
			for i, d := range value.dims {
				dims[i] = C.int64_t(d)
			}
			dimsp = &dims[0]
		}
		C.TF_SetAttrShape(cdesc, cAttrName, dimsp, ndims)
	case []Shape:
		if len(value) == 0 {
			C.TF_SetAttrShapeList(cdesc, cAttrName, nil, nil, 0)
		} else {
			var flatDims []C.int64_t
			ndims := make([]C.int, len(value))
			for i, s := range value {
				nd := s.NumDimensions()
				ndims[i] = C.int(nd)
				for _, d := range s.dims {
					flatDims = append(flatDims, C.int64_t(d))
				}
			}
			var flatDimsp *C.int64_t
			if len(flatDims) > 0 {
				flatDimsp = &flatDims[0]
			}
			C.TF_SetAttrShapeList_Helper(cdesc, cAttrName, flatDimsp, &ndims[0], C.int(len(value)))
		}
	default:
		return fmt.Errorf("attribute %q has a type (%T) which is not valid for operation attributes", name, value)
	}
	return nil
}
<commit_msg>Added go.ImportWithOptions() function and use it to set prefix and default execution device<commit_after>/*
Copyright 2016 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package tensorflow

// #include "tensorflow/c/c_api.h"
//
// #include <stdlib.h>
// #include <string.h>
//
// void TF_SetAttrShapeList_Helper(TF_OperationDescription* desc,
//                                 const char* attr_name,
//                                 const int64_t* flat_dims,
//                                 const int* num_dims,
//                                 int num_shapes) {
//  const int64_t** dims =
//    (const int64_t**)malloc(sizeof(const int64_t*) * num_shapes);
//  int i = 0;
//  for (i = 0; i < num_shapes; i++) {
//    dims[i] = flat_dims;
//    if (num_dims[i] > 0) {
//      // flat_dims will be NULL iff num_shapes is 0 or all elements in num_dims are <= 0.
//      flat_dims += num_dims[i];
//    }
//  }
//  TF_SetAttrShapeList(desc, attr_name, dims, num_dims, num_shapes);
//  free(dims);
// }
import "C"

import (
	"fmt"
	"io"
	"runtime"
	"unsafe"
)

// Graph represents a computation graph. Graphs may be shared between sessions.
type Graph struct {
	c *C.TF_Graph
}

// Graph execution options
type GraphImportOptions struct {
	// Node prefix
	Prefix string

	// Execution device
	Device string

	// TODO: extend this structure to support more options from TF_ImportGraphDefOptions
}

// NewGraph returns a new Graph.
func NewGraph() *Graph {
	g := &Graph{C.TF_NewGraph()}
	runtime.SetFinalizer(g, (*Graph).finalizer)
	return g
}

func (g *Graph) finalizer() {
	C.TF_DeleteGraph(g.c)
}

// WriteTo writes out a serialized representation of g to w.
//
// Implements the io.WriterTo interface.
func (g *Graph) WriteTo(w io.Writer) (int64, error) {
	buf := C.TF_NewBuffer()
	defer C.TF_DeleteBuffer(buf)
	status := newStatus()
	C.TF_GraphToGraphDef(g.c, buf, status.c)
	if err := status.Err(); err != nil {
		return 0, err
	}
	if buf.length > (1 << 30) {
		// For very large graphs, the writes can be chunked.
		// Punt on that for now.
		return 0, fmt.Errorf("Graph is too large to write out, Graph.WriteTo needs to be updated")
	}
	// A []byte slice backed by C memory.
	// See: https://github.com/golang/go/wiki/cgo#turning-c-arrays-into-go-slices
	length := int(buf.length)
	slice := (*[1 << 30]byte)(unsafe.Pointer(buf.data))[:length:length]
	n, err := w.Write(slice)
	return int64(n), err
}

// ImportWithOptions imports the nodes and edges from a serialized representation of
// another Graph into g.
//
// Multiple options can be specified for the newly imported nodes.
func (g *Graph) ImportWithOptions(def []byte, options GraphImportOptions) error {
	cprefix := C.CString(options.Prefix)
	defer C.free(unsafe.Pointer(cprefix))

	opts := C.TF_NewImportGraphDefOptions()
	defer C.TF_DeleteImportGraphDefOptions(opts)
	C.TF_ImportGraphDefOptionsSetPrefix(opts, cprefix)

	if len(options.Device) != 0 {
		cdev := C.CString(options.Device)
		defer C.free(unsafe.Pointer(cdev))
		C.TF_ImportGraphDefOptionsSetDefaultDevice(opts, cdev)
	}

	buf := C.TF_NewBuffer()
	defer C.TF_DeleteBuffer(buf)
	// Would have preferred to use C.CBytes, but that does not play well
	// with "go vet" till https://github.com/golang/go/issues/17201 is
	// resolved.
	buf.length = C.size_t(len(def))
	buf.data = C.malloc(buf.length)
	if buf.data == nil {
		return fmt.Errorf("unable to allocate memory")
	}
	defer C.free(buf.data)
	C.memcpy(buf.data, unsafe.Pointer(&def[0]), buf.length)

	status := newStatus()

	C.TF_GraphImportGraphDef(g.c, buf, opts, status.c)
	if err := status.Err(); err != nil {
		return err
	}

	return nil
}

// Import imports the nodes and edges from a serialized representation of
// another Graph into g.
//
// Names of imported nodes will be prefixed with prefix.
func (g *Graph) Import(def []byte, prefix string) error {
    return g.ImportWithOptions(def, GraphImportOptions{Prefix: prefix})
}

// Operation returns the Operation named name in the Graph, or nil if no such
// operation is present.
func (g *Graph) Operation(name string) *Operation {
	cname := C.CString(name)
	defer C.free(unsafe.Pointer(cname))
	cop := C.TF_GraphOperationByName(g.c, cname)
	if cop == nil {
		return nil
	}
	return &Operation{cop, g}
}

// Operations returns a list of all operations in the graph
func (g *Graph) Operations() []Operation {
	var pos C.size_t = 0
	ops := []Operation{}
	for {
		cop := C.TF_GraphNextOperation(g.c, &pos)
		if cop == nil {
			break
		}
		ops = append(ops, Operation{cop, g})
	}
	return ops
}

// OpSpec is the specification of an Operation to be added to a Graph
// (using Graph.AddOperation).
type OpSpec struct {
	// Type of the operation (e.g., "Add", "MatMul").
	Type string

	// Name by which the added operation will be referred to in the Graph.
	// If omitted, defaults to Type.
	Name string

	// Inputs to this operation, which in turn must be outputs
	// of other operations already added to the Graph.
	//
	// An operation may have multiple inputs with individual inputs being
	// either a single tensor produced by another operation or a list of
	// tensors produced by multiple operations. For example, the "Concat"
	// operation takes two inputs: (1) the dimension along which to
	// concatenate and (2) a list of tensors to concatenate. Thus, for
	// Concat, len(Input) must be 2, with the first element being an Output
	// and the second being an OutputList.
	Input []Input

	// Map from attribute name to its value that will be attached to this
	// operation.
	Attrs map[string]interface{}

	// Operations that must be executed before executing the operation
	// being added.
	ControlDependencies []*Operation

	// The device on which the operation should be executed.
	// If omitted, an appropriate device will automatically be selected.
	//
	// For example, if set of "/device:GPU:0", then the operation will
	// execute on GPU #0.
	Device string

	// Other possible fields: ColocateWith.
}

// AddOperation adds an operation to g.
func (g *Graph) AddOperation(args OpSpec) (*Operation, error) {
	if args.Name == "" {
		args.Name = args.Type
	}
	cname := C.CString(args.Name)
	ctype := C.CString(args.Type)
	cdesc := C.TF_NewOperation(g.c, ctype, cname)
	C.free(unsafe.Pointer(cname))
	C.free(unsafe.Pointer(ctype))

	for _, in := range args.Input {
		switch in := in.(type) {
		case Output:
			C.TF_AddInput(cdesc, in.c())
		case OutputList:
			size := len(in)
			list := make([]C.TF_Output, size)
			for i, v := range in {
				list[i] = v.c()
			}
			if size > 0 {
				C.TF_AddInputList(cdesc, &list[0], C.int(size))
			} else {
				C.TF_AddInputList(cdesc, nil, 0)
			}
		}
	}
	for _, in := range args.ControlDependencies {
		C.TF_AddControlInput(cdesc, in.c)
	}
	status := newStatus()
	for name, value := range args.Attrs {
		if err := setAttr(cdesc, status, name, value); err != nil {
			// Memory leak here as the TF_OperationDescription
			// object will not be cleaned up. At the time of this
			// writing, this was next to impossible since it
			// required value to be a string tensor with
			// incorrectly encoded strings. Given this rarity, live
			// with the memory leak.  If it becomes a real problem,
			// consider adding a TF_DeleteOperationDescription
			// function to the C API.
			return nil, fmt.Errorf("%v (memory will be leaked)", err)
		}
	}
	if len(args.Device) > 0 {
		cdevice := C.CString(args.Device)
		C.TF_SetDevice(cdesc, cdevice)
		C.free(unsafe.Pointer(cdevice))
	}
	c := C.TF_FinishOperation(cdesc, status.c)
	if err := status.Err(); err != nil {
		return nil, err
	}
	return &Operation{c, g}, nil
}

func setAttr(cdesc *C.TF_OperationDescription, status *status, name string, value interface{}) error {
	cAttrName := C.CString(name)
	defer C.free(unsafe.Pointer(cAttrName))
	switch value := value.(type) {
	case string:
		cstr := C.CString(value)
		C.TF_SetAttrString(cdesc, cAttrName, unsafe.Pointer(cstr), C.size_t(len(value)))
		C.free(unsafe.Pointer(cstr))
	case []string:
		size := len(value)
		list := make([]unsafe.Pointer, size)
		lens := make([]C.size_t, size)
		for i, s := range value {
			list[i] = unsafe.Pointer(C.CString(s))
			lens[i] = C.size_t(len(s))
		}
		if size > 0 {
			C.TF_SetAttrStringList(cdesc, cAttrName, &list[0], &lens[0], C.int(size))
		} else {
			C.TF_SetAttrStringList(cdesc, cAttrName, nil, nil, 0)
		}
		for _, s := range list {
			C.free(s)
		}
	case int64:
		C.TF_SetAttrInt(cdesc, cAttrName, C.int64_t(value))
	case []int64:
		size := len(value)
		list := make([]C.int64_t, size)
		for i, v := range value {
			list[i] = C.int64_t(v)
		}
		if size > 0 {
			C.TF_SetAttrIntList(cdesc, cAttrName, &list[0], C.int(size))
		} else {
			C.TF_SetAttrIntList(cdesc, cAttrName, nil, 0)
		}
	case float32:
		C.TF_SetAttrFloat(cdesc, cAttrName, C.float(value))
	case []float32:
		size := len(value)
		list := make([]C.float, size)
		for i, v := range value {
			list[i] = C.float(v)
		}
		if size > 0 {
			C.TF_SetAttrFloatList(cdesc, cAttrName, &list[0], C.int(size))
		} else {
			C.TF_SetAttrFloatList(cdesc, cAttrName, nil, 0)
		}
	case bool:
		v := C.uchar(0)
		if value {
			v = 1
		}
		C.TF_SetAttrBool(cdesc, cAttrName, v)
	case []bool:
		size := len(value)
		list := make([]C.uchar, size)
		for i, v := range value {
			if v {
				list[i] = 1
			}
		}
		if size > 0 {
			C.TF_SetAttrBoolList(cdesc, cAttrName, &list[0], C.int(size))
		} else {
			C.TF_SetAttrBoolList(cdesc, cAttrName, nil, 0)
		}
	case DataType:
		C.TF_SetAttrType(cdesc, cAttrName, C.TF_DataType(value))
	case []DataType:
		var list *C.TF_DataType
		if len(value) > 0 {
			list = (*C.TF_DataType)(&value[0])
		}
		C.TF_SetAttrTypeList(cdesc, cAttrName, list, C.int(len(value)))
	case *Tensor:
		C.TF_SetAttrTensor(cdesc, cAttrName, value.c, status.c)
		if err := status.Err(); err != nil {
			return fmt.Errorf("bad value for attribute %q: %v", name, err)
		}
	case []*Tensor:
		size := len(value)
		list := make([]*C.TF_Tensor, size)
		for i, v := range value {
			list[i] = v.c
		}
		var plist **C.TF_Tensor
		if size > 0 {
			plist = &list[0]
		}
		C.TF_SetAttrTensorList(cdesc, cAttrName, plist, C.int(size), status.c)
		if err := status.Err(); err != nil {
			return fmt.Errorf("bad value for attribute %q: %v", name, err)
		}
	case Shape:
		ndims := C.int(value.NumDimensions())
		var dimsp *C.int64_t
		if ndims > 0 {
			dims := make([]C.int64_t, ndims)
			for i, d := range value.dims {
				dims[i] = C.int64_t(d)
			}
			dimsp = &dims[0]
		}
		C.TF_SetAttrShape(cdesc, cAttrName, dimsp, ndims)
	case []Shape:
		if len(value) == 0 {
			C.TF_SetAttrShapeList(cdesc, cAttrName, nil, nil, 0)
		} else {
			var flatDims []C.int64_t
			ndims := make([]C.int, len(value))
			for i, s := range value {
				nd := s.NumDimensions()
				ndims[i] = C.int(nd)
				for _, d := range s.dims {
					flatDims = append(flatDims, C.int64_t(d))
				}
			}
			var flatDimsp *C.int64_t
			if len(flatDims) > 0 {
				flatDimsp = &flatDims[0]
			}
			C.TF_SetAttrShapeList_Helper(cdesc, cAttrName, flatDimsp, &ndims[0], C.int(len(value)))
		}
	default:
		return fmt.Errorf("attribute %q has a type (%T) which is not valid for operation attributes", name, value)
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2009 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.

// TODO(cw): ListenPacket test, Read() test, ipv6 test &
// Dial()/Listen() level tests

package net

import (
	"bytes"
	"flag"
	"os"
	"testing"
)

const ICMP_ECHO_REQUEST = 8
const ICMP_ECHO_REPLY = 0

// returns a suitable 'ping request' packet, with id & seq and a
// payload length of pktlen
func makePingRequest(id, seq, pktlen int, filler []byte) []byte {
	p := make([]byte, pktlen)
	copy(p[8:], bytes.Repeat(filler, (pktlen-8)/len(filler)+1))

	p[0] = ICMP_ECHO_REQUEST // type
	p[1] = 0                 // code
	p[2] = 0                 // cksum
	p[3] = 0                 // cksum
	p[4] = uint8(id >> 8)    // id
	p[5] = uint8(id & 0xff)  // id
	p[6] = uint8(seq >> 8)   // sequence
	p[7] = uint8(seq & 0xff) // sequence

	// calculate icmp checksum
	cklen := len(p)
	s := uint32(0)
	for i := 0; i < (cklen - 1); i += 2 {
		s += uint32(p[i+1])<<8 | uint32(p[i])
	}
	if cklen&1 == 1 {
		s += uint32(p[cklen-1])
	}
	s = (s >> 16) + (s & 0xffff)
	s = s + (s >> 16)

	// place checksum back in header; using ^= avoids the
	// assumption the checksum bytes are zero
	p[2] ^= uint8(^s & 0xff)
	p[3] ^= uint8(^s >> 8)

	return p
}

func parsePingReply(p []byte) (id, seq int) {
	id = int(p[4])<<8 | int(p[5])
	seq = int(p[6])<<8 | int(p[7])
	return
}

var srchost = flag.String("srchost", "", "Source of the ICMP ECHO request")
var dsthost = flag.String("dsthost", "localhost", "Destination for the ICMP ECHO request")

// test (raw) IP socket using ICMP
func TestICMP(t *testing.T) {
	if os.Getuid() != 0 {
		t.Logf("test disabled; must be root")
		return
	}

	var laddr *IPAddr
	if *srchost != "" {
		laddr, err := ResolveIPAddr(*srchost)
		if err != nil {
			t.Fatalf(`net.ResolveIPAddr("%v") = %v, %v`, *srchost, laddr, err)
		}
	}

	raddr, err := ResolveIPAddr(*dsthost)
	if err != nil {
		t.Fatalf(`net.ResolveIPAddr("%v") = %v, %v`, *dsthost, raddr, err)
	}

	c, err := ListenIP("ip4:icmp", laddr)
	if err != nil {
		t.Fatalf(`net.ListenIP("ip4:icmp", %v) = %v, %v`, *srchost, c, err)
	}

	sendid := os.Getpid()
	const sendseq = 61455
	const pingpktlen = 128
	sendpkt := makePingRequest(sendid, sendseq, pingpktlen, []byte("Go Go Gadget Ping!!!"))

	n, err := c.WriteToIP(sendpkt, raddr)
	if err != nil || n != pingpktlen {
		t.Fatalf(`net.WriteToIP(..., %v) = %v, %v`, raddr, n, err)
	}

	c.SetTimeout(100e6)
	resp := make([]byte, 1024)
	for {
		n, from, err := c.ReadFrom(resp)
		if err != nil {
			t.Fatalf(`ReadFrom(...) = %v, %v, %v`, n, from, err)
		}
		if resp[0] != ICMP_ECHO_REPLY {
			continue
		}
		rcvid, rcvseq := parsePingReply(resp)
		if rcvid != sendid || rcvseq != sendseq {
			t.Fatal(`Ping reply saw id,seq=%v,%v (expected %v, %v)`, rcvid, rcvseq, sendid, sendseq)
		}
		return
	}
	t.Fatalf("saw no ping return")
}
<commit_msg>net: Fix ICMP test.<commit_after>// Copyright 2009 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.

// TODO(cw): ListenPacket test, Read() test, ipv6 test &
// Dial()/Listen() level tests

package net

import (
	"bytes"
	"flag"
	"os"
	"testing"
)

const ICMP_ECHO_REQUEST = 8
const ICMP_ECHO_REPLY = 0

// returns a suitable 'ping request' packet, with id & seq and a
// payload length of pktlen
func makePingRequest(id, seq, pktlen int, filler []byte) []byte {
	p := make([]byte, pktlen)
	copy(p[8:], bytes.Repeat(filler, (pktlen-8)/len(filler)+1))

	p[0] = ICMP_ECHO_REQUEST // type
	p[1] = 0                 // code
	p[2] = 0                 // cksum
	p[3] = 0                 // cksum
	p[4] = uint8(id >> 8)    // id
	p[5] = uint8(id & 0xff)  // id
	p[6] = uint8(seq >> 8)   // sequence
	p[7] = uint8(seq & 0xff) // sequence

	// calculate icmp checksum
	cklen := len(p)
	s := uint32(0)
	for i := 0; i < (cklen - 1); i += 2 {
		s += uint32(p[i+1])<<8 | uint32(p[i])
	}
	if cklen&1 == 1 {
		s += uint32(p[cklen-1])
	}
	s = (s >> 16) + (s & 0xffff)
	s = s + (s >> 16)

	// place checksum back in header; using ^= avoids the
	// assumption the checksum bytes are zero
	p[2] ^= uint8(^s & 0xff)
	p[3] ^= uint8(^s >> 8)

	return p
}

func parsePingReply(p []byte) (id, seq int) {
	id = int(p[4])<<8 | int(p[5])
	seq = int(p[6])<<8 | int(p[7])
	return
}

var srchost = flag.String("srchost", "", "Source of the ICMP ECHO request")
var dsthost = flag.String("dsthost", "localhost", "Destination for the ICMP ECHO request")

// test (raw) IP socket using ICMP
func TestICMP(t *testing.T) {
	if os.Getuid() != 0 {
		t.Logf("test disabled; must be root")
		return
	}

	var laddr *IPAddr
	if *srchost != "" {
		laddr, err := ResolveIPAddr(*srchost)
		if err != nil {
			t.Fatalf(`net.ResolveIPAddr("%v") = %v, %v`, *srchost, laddr, err)
		}
	}

	raddr, err := ResolveIPAddr(*dsthost)
	if err != nil {
		t.Fatalf(`net.ResolveIPAddr("%v") = %v, %v`, *dsthost, raddr, err)
	}

	c, err := ListenIP("ip4:icmp", laddr)
	if err != nil {
		t.Fatalf(`net.ListenIP("ip4:icmp", %v) = %v, %v`, *srchost, c, err)
	}

	sendid := os.Getpid() & 0xffff
	const sendseq = 61455
	const pingpktlen = 128
	sendpkt := makePingRequest(sendid, sendseq, pingpktlen, []byte("Go Go Gadget Ping!!!"))

	n, err := c.WriteToIP(sendpkt, raddr)
	if err != nil || n != pingpktlen {
		t.Fatalf(`net.WriteToIP(..., %v) = %v, %v`, raddr, n, err)
	}

	c.SetTimeout(100e6)
	resp := make([]byte, 1024)
	for {
		n, from, err := c.ReadFrom(resp)
		if err != nil {
			t.Fatalf(`ReadFrom(...) = %v, %v, %v`, n, from, err)
		}
		if resp[0] != ICMP_ECHO_REPLY {
			continue
		}
		rcvid, rcvseq := parsePingReply(resp)
		if rcvid != sendid || rcvseq != sendseq {
			t.Fatalf(`Ping reply saw id,seq=0x%x,0x%x (expected 0x%x, 0x%x)`, rcvid, rcvseq, sendid, sendseq)
		}
		return
	}
	t.Fatalf("saw no ping return")
}
<|endoftext|>


<commit_before>// Copyright 2009 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 runtime contains operations that interact with Go's runtime system,
	such as functions to control goroutines. It also includes the low-level type information
	used by the reflect package; see reflect's documentation for the programmable
	interface to the run-time type system.
*/
package runtime

// Gosched yields the processor, allowing other goroutines to run.  It does not
// suspend the current goroutine, so execution resumes automatically.
func Gosched()

// Goexit terminates the goroutine that calls it.  No other goroutine is affected.
// Goexit runs all deferred calls before terminating the goroutine.
func Goexit()

// Caller reports file and line number information about function invocations on
// the calling goroutine's stack.  The argument skip is the number of stack frames
// to ascend, with 0 identifying the caller of Caller.  The return values report the
// program counter, file name, and line number within the file of the corresponding
// call.  The boolean ok is false if it was not possible to recover the information.
func Caller(skip int) (pc uintptr, file string, line int, ok bool)

// Callers fills the slice pc with the program counters of function invocations
// on the calling goroutine's stack.  The argument skip is the number of stack frames
// to skip before recording in pc, with 0 starting at the caller of Caller.
// It returns the number of entries written to pc.
func Callers(skip int, pc []uintptr) int

type Func struct { // Keep in sync with runtime.h:struct Func
	name   string
	typ    string  // go type string
	src    string  // src file name
	pcln   []byte  // pc/ln tab for this func
	entry  uintptr // entry pc
	pc0    uintptr // starting pc, ln for table
	ln0    int32
	frame  int32 // stack frame size
	args   int32 // number of 32-bit in/out args
	locals int32 // number of 32-bit locals
}

// FuncForPC returns a *Func describing the function that contains the
// given program counter address, or else nil.
func FuncForPC(pc uintptr) *Func

// Name returns the name of the function.
func (f *Func) Name() string { return f.name }

// Entry returns the entry address of the function.
func (f *Func) Entry() uintptr { return f.entry }

// FileLine returns the file name and line number of the
// source code corresponding to the program counter pc.
// The result will not be accurate if pc is not a program
// counter within f.
func (f *Func) FileLine(pc uintptr) (file string, line int) {
	return funcline_go(f, pc)
}

// implemented in symtab.c
func funcline_go(*Func, uintptr) (string, int)

// mid returns the current os thread (m) id.
func mid() uint32

// NumCPU returns the number of logical CPUs on the local machine.
func NumCPU() int

// Semacquire waits until *s > 0 and then atomically decrements it.
// It is intended as a simple sleep primitive for use by the synchronization
// library and should not be used directly.
func Semacquire(s *uint32)

// Semrelease atomically increments *s and notifies a waiting goroutine
// if one is blocked in Semacquire.
// It is intended as a simple wakeup primitive for use by the synchronization
// library and should not be used directly.
func Semrelease(s *uint32)

// SetFinalizer sets the finalizer associated with x to f.
// When the garbage collector finds an unreachable block
// with an associated finalizer, it clears the association and runs
// f(x) in a separate goroutine.  This makes x reachable again, but
// now without an associated finalizer.  Assuming that SetFinalizer
// is not called again, the next time the garbage collector sees
// that x is unreachable, it will free x.
//
// SetFinalizer(x, nil) clears any finalizer associated with x.
//
// The argument x must be a pointer to an object allocated by
// calling new or by taking the address of a composite literal.
// The argument f must be a function that takes a single argument
// of x's type and can have arbitrary ignored return values.
// If either of these is not true, SetFinalizer aborts the program.
//
// Finalizers are run in dependency order: if A points at B, both have
// finalizers, and they are otherwise unreachable, only the finalizer
// for A runs; once A is freed, the finalizer for B can run.
// If a cyclic structure includes a block with a finalizer, that
// cycle is not guaranteed to be garbage collected and the finalizer
// is not guaranteed to run, because there is no ordering that
// respects the dependencies.
//
// The finalizer for x is scheduled to run at some arbitrary time after
// x becomes unreachable.
// There is no guarantee that finalizers will run before a program exits,
// so typically they are useful only for releasing non-memory resources
// associated with an object during a long-running program.
// For example, an os.File object could use a finalizer to close the
// associated operating system file descriptor when a program discards
// an os.File without calling Close, but it would be a mistake
// to depend on a finalizer to flush an in-memory I/O buffer such as a
// bufio.Writer, because the buffer would not be flushed at program exit.
//
// A single goroutine runs all finalizers for a program, sequentially.
// If a finalizer must run for a long time, it should do so by starting
// a new goroutine.
func SetFinalizer(x, f interface{})

func getgoroot() string

// GOROOT returns the root of the Go tree.
// It uses the GOROOT environment variable, if set,
// or else the root used during the Go build.
func GOROOT() string {
	s := getgoroot()
	if s != "" {
		return s
	}
	return defaultGoroot
}

// Version returns the Go tree's version string.
// It is either a sequence number or, when possible,
// a release tag like "release.2010-03-04".
// A trailing + indicates that the tree had local modifications
// at the time of the build.
func Version() string {
	return theVersion
}

// GOOS is the Go tree's operating system target:
// one of darwin, freebsd, linux, and so on.
const GOOS string = theGoos

// GOARCH is the Go tree's architecture target:
// 386, amd64, or arm.
const GOARCH string = theGoarch
<commit_msg>runtime: fix typo in comment<commit_after>// Copyright 2009 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 runtime contains operations that interact with Go's runtime system,
	such as functions to control goroutines. It also includes the low-level type information
	used by the reflect package; see reflect's documentation for the programmable
	interface to the run-time type system.
*/
package runtime

// Gosched yields the processor, allowing other goroutines to run.  It does not
// suspend the current goroutine, so execution resumes automatically.
func Gosched()

// Goexit terminates the goroutine that calls it.  No other goroutine is affected.
// Goexit runs all deferred calls before terminating the goroutine.
func Goexit()

// Caller reports file and line number information about function invocations on
// the calling goroutine's stack.  The argument skip is the number of stack frames
// to ascend, with 0 identifying the caller of Caller.  The return values report the
// program counter, file name, and line number within the file of the corresponding
// call.  The boolean ok is false if it was not possible to recover the information.
func Caller(skip int) (pc uintptr, file string, line int, ok bool)

// Callers fills the slice pc with the program counters of function invocations
// on the calling goroutine's stack.  The argument skip is the number of stack frames
// to skip before recording in pc, with 0 starting at the caller of Callers.
// It returns the number of entries written to pc.
func Callers(skip int, pc []uintptr) int

type Func struct { // Keep in sync with runtime.h:struct Func
	name   string
	typ    string  // go type string
	src    string  // src file name
	pcln   []byte  // pc/ln tab for this func
	entry  uintptr // entry pc
	pc0    uintptr // starting pc, ln for table
	ln0    int32
	frame  int32 // stack frame size
	args   int32 // number of 32-bit in/out args
	locals int32 // number of 32-bit locals
}

// FuncForPC returns a *Func describing the function that contains the
// given program counter address, or else nil.
func FuncForPC(pc uintptr) *Func

// Name returns the name of the function.
func (f *Func) Name() string { return f.name }

// Entry returns the entry address of the function.
func (f *Func) Entry() uintptr { return f.entry }

// FileLine returns the file name and line number of the
// source code corresponding to the program counter pc.
// The result will not be accurate if pc is not a program
// counter within f.
func (f *Func) FileLine(pc uintptr) (file string, line int) {
	return funcline_go(f, pc)
}

// implemented in symtab.c
func funcline_go(*Func, uintptr) (string, int)

// mid returns the current os thread (m) id.
func mid() uint32

// NumCPU returns the number of logical CPUs on the local machine.
func NumCPU() int

// Semacquire waits until *s > 0 and then atomically decrements it.
// It is intended as a simple sleep primitive for use by the synchronization
// library and should not be used directly.
func Semacquire(s *uint32)

// Semrelease atomically increments *s and notifies a waiting goroutine
// if one is blocked in Semacquire.
// It is intended as a simple wakeup primitive for use by the synchronization
// library and should not be used directly.
func Semrelease(s *uint32)

// SetFinalizer sets the finalizer associated with x to f.
// When the garbage collector finds an unreachable block
// with an associated finalizer, it clears the association and runs
// f(x) in a separate goroutine.  This makes x reachable again, but
// now without an associated finalizer.  Assuming that SetFinalizer
// is not called again, the next time the garbage collector sees
// that x is unreachable, it will free x.
//
// SetFinalizer(x, nil) clears any finalizer associated with x.
//
// The argument x must be a pointer to an object allocated by
// calling new or by taking the address of a composite literal.
// The argument f must be a function that takes a single argument
// of x's type and can have arbitrary ignored return values.
// If either of these is not true, SetFinalizer aborts the program.
//
// Finalizers are run in dependency order: if A points at B, both have
// finalizers, and they are otherwise unreachable, only the finalizer
// for A runs; once A is freed, the finalizer for B can run.
// If a cyclic structure includes a block with a finalizer, that
// cycle is not guaranteed to be garbage collected and the finalizer
// is not guaranteed to run, because there is no ordering that
// respects the dependencies.
//
// The finalizer for x is scheduled to run at some arbitrary time after
// x becomes unreachable.
// There is no guarantee that finalizers will run before a program exits,
// so typically they are useful only for releasing non-memory resources
// associated with an object during a long-running program.
// For example, an os.File object could use a finalizer to close the
// associated operating system file descriptor when a program discards
// an os.File without calling Close, but it would be a mistake
// to depend on a finalizer to flush an in-memory I/O buffer such as a
// bufio.Writer, because the buffer would not be flushed at program exit.
//
// A single goroutine runs all finalizers for a program, sequentially.
// If a finalizer must run for a long time, it should do so by starting
// a new goroutine.
func SetFinalizer(x, f interface{})

func getgoroot() string

// GOROOT returns the root of the Go tree.
// It uses the GOROOT environment variable, if set,
// or else the root used during the Go build.
func GOROOT() string {
	s := getgoroot()
	if s != "" {
		return s
	}
	return defaultGoroot
}

// Version returns the Go tree's version string.
// It is either a sequence number or, when possible,
// a release tag like "release.2010-03-04".
// A trailing + indicates that the tree had local modifications
// at the time of the build.
func Version() string {
	return theVersion
}

// GOOS is the Go tree's operating system target:
// one of darwin, freebsd, linux, and so on.
const GOOS string = theGoos

// GOARCH is the Go tree's architecture target:
// 386, amd64, or arm.
const GOARCH string = theGoarch
<|endoftext|>


<commit_before>// Copyright 2009 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 strings

import (
	"errors"
	"io"
	"unicode/utf8"
)

// A Reader implements the io.Reader, io.Seeker, io.ByteScanner, and
// io.RuneScanner interfaces by reading from a string.
type Reader struct {
	s        string
	i        int // current reading index
	prevRune int // index of previous rune; or < 0
}

// Len returns the number of bytes of the unread portion of the
// string.
func (r *Reader) Len() int {
	if r.i >= len(r.s) {
		return 0
	}
	return len(r.s) - r.i
}

func (r *Reader) Read(b []byte) (n int, err error) {
	if len(b) == 0 {
		return 0, nil
	}
	if r.i >= len(r.s) {
		return 0, io.EOF
	}
	n = copy(b, r.s[r.i:])
	r.i += n
	r.prevRune = -1
	return
}

func (r *Reader) ReadByte() (b byte, err error) {
	if r.i >= len(r.s) {
		return 0, io.EOF
	}
	b = r.s[r.i]
	r.i++
	r.prevRune = -1
	return
}

// UnreadByte moves the reading position back by one byte.
// It is an error to call UnreadByte if nothing has been
// read yet.
func (r *Reader) UnreadByte() error {
	if r.i <= 0 {
		return errors.New("strings.Reader: at beginning of string")
	}
	r.i--
	r.prevRune = -1
	return nil
}

// ReadRune reads and returns the next UTF-8-encoded
// Unicode code point from the buffer.
// If no bytes are available, the error returned is io.EOF.
// If the bytes are an erroneous UTF-8 encoding, it
// consumes one byte and returns U+FFFD, 1.
func (r *Reader) ReadRune() (ch rune, size int, err error) {
	if r.i >= len(r.s) {
		return 0, 0, io.EOF
	}
	r.prevRune = r.i
	if c := r.s[r.i]; c < utf8.RuneSelf {
		r.i++
		return rune(c), 1, nil
	}
	ch, size = utf8.DecodeRuneInString(r.s[r.i:])
	r.i += size
	return
}

// UnreadRune causes the next call to ReadRune to return the same rune
// as the previous call to ReadRune.
// The last method called on r must have been ReadRune.
func (r *Reader) UnreadRune() error {
	if r.prevRune < 0 {
		return errors.New("strings.Reader: previous operation was not ReadRune")
	}
	r.i = r.prevRune
	r.prevRune = -1
	return nil
}

// Seek implements the io.Seeker interface.
func (r *Reader) Seek(offset int64, whence int) (int64, error) {
	var abs int64
	switch whence {
	case 0:
		abs = offset
	case 1:
		abs = int64(r.i) + offset
	case 2:
		abs = int64(len(r.s)) + offset
	default:
		return 0, errors.New("strings: invalid whence")
	}
	if abs < 0 {
		return 0, errors.New("strings: negative position")
	}
	if abs >= 1<<31 {
		return 0, errors.New("strings: position out of range")
	}
	r.i = int(abs)
	return abs, nil
}

// NewReader returns a new Reader reading from s.
// It is similar to bytes.NewBufferString but more efficient and read-only.
func NewReader(s string) *Reader { return &Reader{s, 0, -1} }
<commit_msg>strings: delete method comments implied by interface satisfaction<commit_after>// Copyright 2009 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 strings

import (
	"errors"
	"io"
	"unicode/utf8"
)

// A Reader implements the io.Reader, io.Seeker, io.ByteScanner, and
// io.RuneScanner interfaces by reading from a string.
type Reader struct {
	s        string
	i        int // current reading index
	prevRune int // index of previous rune; or < 0
}

// Len returns the number of bytes of the unread portion of the
// string.
func (r *Reader) Len() int {
	if r.i >= len(r.s) {
		return 0
	}
	return len(r.s) - r.i
}

func (r *Reader) Read(b []byte) (n int, err error) {
	if len(b) == 0 {
		return 0, nil
	}
	if r.i >= len(r.s) {
		return 0, io.EOF
	}
	n = copy(b, r.s[r.i:])
	r.i += n
	r.prevRune = -1
	return
}

func (r *Reader) ReadByte() (b byte, err error) {
	if r.i >= len(r.s) {
		return 0, io.EOF
	}
	b = r.s[r.i]
	r.i++
	r.prevRune = -1
	return
}

func (r *Reader) UnreadByte() error {
	if r.i <= 0 {
		return errors.New("strings.Reader: at beginning of string")
	}
	r.i--
	r.prevRune = -1
	return nil
}

func (r *Reader) ReadRune() (ch rune, size int, err error) {
	if r.i >= len(r.s) {
		return 0, 0, io.EOF
	}
	r.prevRune = r.i
	if c := r.s[r.i]; c < utf8.RuneSelf {
		r.i++
		return rune(c), 1, nil
	}
	ch, size = utf8.DecodeRuneInString(r.s[r.i:])
	r.i += size
	return
}

func (r *Reader) UnreadRune() error {
	if r.prevRune < 0 {
		return errors.New("strings.Reader: previous operation was not ReadRune")
	}
	r.i = r.prevRune
	r.prevRune = -1
	return nil
}

// Seek implements the io.Seeker interface.
func (r *Reader) Seek(offset int64, whence int) (int64, error) {
	var abs int64
	switch whence {
	case 0:
		abs = offset
	case 1:
		abs = int64(r.i) + offset
	case 2:
		abs = int64(len(r.s)) + offset
	default:
		return 0, errors.New("strings: invalid whence")
	}
	if abs < 0 {
		return 0, errors.New("strings: negative position")
	}
	if abs >= 1<<31 {
		return 0, errors.New("strings: position out of range")
	}
	r.i = int(abs)
	return abs, nil
}

// NewReader returns a new Reader reading from s.
// It is similar to bytes.NewBufferString but more efficient and read-only.
func NewReader(s string) *Reader { return &Reader{s, 0, -1} }
<|endoftext|>


<commit_before>package build

import (
	"strings"
	"testing"
)

func TestPropertyToStringString(t *testing.T) {
	expected := `"test"`
	actual, err := PropertyToString("test", true)
	if err != nil || actual != expected {
		t.Error("Error serializing string", err)
	}
}

func TestPropertyToStringInt(t *testing.T) {
	expected := `42`
	actual, err := PropertyToString(42, true)
	if err != nil || actual != expected {
		t.Error("Error serializing int", err)
	}
}

func TestPropertyToStringFloat(t *testing.T) {
	expected := `4.2`
	actual, err := PropertyToString(4.2, true)
	if err != nil || actual != expected {
		t.Error("Error serializing float", err)
	}
}

func TestPropertyToStringList(t *testing.T) {
	expected := `[1, 2, 3]`
	actual, err := PropertyToString([]int{1, 2, 3}, true)
	if err != nil || actual != expected {
		t.Error("Error serializing slice", err)
	}
}

func TestPropertyToStringCompositeList(t *testing.T) {
	expected := `[1, 2, 3, "spam"]`
	actual, err := PropertyToString([]interface{}{1, 2, 3, "spam"}, true)
	if err != nil || actual != expected {
		t.Error("Error serializing composite slice", err)
	}
}

func TestPropertyToStringMap(t *testing.T) {
	expected := `["bar": 2, "foo": 1]`
	actual, err := PropertyToString(map[string]int{"foo": 1, "bar": 2}, true)
	if err != nil || actual != expected {
		t.Error("Error serializing map", err)
	}
}

func TestPropertyToStringCompositeMap(t *testing.T) {
	expected := `["bar": 2, "foo": 1, 3: "spam"]`
	actual, err := PropertyToString(map[interface{}]interface{}{"foo": 1, "bar": 2, 3: "spam"}, true)
	if err != nil || actual != expected {
		t.Error("Error serializing composite map", err)
	}
}

func TestGetEnvironmentSimple(t *testing.T) {
	context := &Context{
		Environment: map[string]string{
			"FOO": "BAR",
		},
		Build: &Build{
			Dir: "dir",
		},
	}
	env, err := context.GetEnvironment()
	if err != nil {
		t.Errorf("Error getting environment: %v", err)
	}
	for _, line := range env {
		if line == "FOO=BAR" {
			return
		}
	}
	t.Error("Env FOO=BAR not found")
}

func TestGetEnvironmentComplex(t *testing.T) {
	context := &Context{
		Environment: map[string]string{
			"FOO": "BAR:${HOME}",
		},
		Build: &Build{
			Dir: "dir",
		},
	}
	env, err := context.GetEnvironment()
	if err != nil {
		t.Errorf("Error getting environment: %v", err)
	}
	var foo string
	for _, line := range env {
		if strings.HasPrefix(line, "HOME=") {
			foo = "FOO=BAR:" + line[5:]
		}
	}
	if foo == "" {
		return
	}
	for _, line := range env {
		if line == foo {
			return
		}
	}
	t.Error("Environment variable FOO not set correctly")
}
<commit_msg>Fixed null property values<commit_after>package build

import (
	"strings"
	"testing"
)

func TestPropertyToStringString(t *testing.T) {
	expected := `"test"`
	actual, err := PropertyToString("test", true)
	if err != nil || actual != expected {
		t.Error("Error serializing string", err)
	}
}

func TestPropertyToStringInt(t *testing.T) {
	expected := `42`
	actual, err := PropertyToString(42, true)
	if err != nil || actual != expected {
		t.Error("Error serializing int", err)
	}
}

func TestPropertyToStringFloat(t *testing.T) {
	expected := `4.2`
	actual, err := PropertyToString(4.2, true)
	if err != nil || actual != expected {
		t.Error("Error serializing float", err)
	}
}

func TestPropertyToStringList(t *testing.T) {
	expected := `[1, 2, 3]`
	actual, err := PropertyToString([]int{1, 2, 3}, true)
	if err != nil || actual != expected {
		t.Error("Error serializing slice", err)
	}
}

func TestPropertyToStringCompositeList(t *testing.T) {
	expected := `[1, 2, 3, "spam"]`
	actual, err := PropertyToString([]interface{}{1, 2, 3, "spam"}, true)
	if err != nil || actual != expected {
		t.Error("Error serializing composite slice", err)
	}
}

func TestPropertyToStringMap(t *testing.T) {
	expected := `["bar": 2, "foo": 1]`
	actual, err := PropertyToString(map[string]int{"foo": 1, "bar": 2}, true)
	if err != nil || actual != expected {
		t.Error("Error serializing map", err)
	}
}

func TestPropertyToStringCompositeMap(t *testing.T) {
	expected := `["bar": 2, "foo": 1, 3: "spam"]`
	actual, err := PropertyToString(map[interface{}]interface{}{"foo": 1, "bar": 2, 3: "spam"}, true)
	if err != nil || actual != expected {
		t.Error("Error serializing composite map", err)
	}
}

func TestPropertyToStringNull(t *testing.T) {
	expected := `null`
	actual, err := PropertyToString(nil, true)
	if err != nil || actual != expected {
		t.Error("Error serializing null", err)
	}
}

func TestGetEnvironmentSimple(t *testing.T) {
	context := &Context{
		Environment: map[string]string{
			"FOO": "BAR",
		},
		Build: &Build{
			Dir: "dir",
		},
	}
	env, err := context.GetEnvironment()
	if err != nil {
		t.Errorf("Error getting environment: %v", err)
	}
	for _, line := range env {
		if line == "FOO=BAR" {
			return
		}
	}
	t.Error("Env FOO=BAR not found")
}

func TestGetEnvironmentComplex(t *testing.T) {
	context := &Context{
		Environment: map[string]string{
			"FOO": "BAR:${HOME}",
		},
		Build: &Build{
			Dir: "dir",
		},
	}
	env, err := context.GetEnvironment()
	if err != nil {
		t.Errorf("Error getting environment: %v", err)
	}
	var foo string
	for _, line := range env {
		if strings.HasPrefix(line, "HOME=") {
			foo = "FOO=BAR:" + line[5:]
		}
	}
	if foo == "" {
		return
	}
	for _, line := range env {
		if line == foo {
			return
		}
	}
	t.Error("Environment variable FOO not set correctly")
}
<|endoftext|>


<commit_before>package build

import (
	"reflect"
	"strings"
	"testing"
)

func TestEvaluateString(t *testing.T) {
	context := NewContext(nil)
	if actual, err := context.EvaluateString(`foo`); err != nil || actual != `foo` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo ={`); err != nil || actual != `foo ={` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo #{"bar"}`); err != nil || actual != `foo bar` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo ={"bar"}`); err != nil || actual != `foo bar` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo ={1+1}`); err != nil || actual != `foo 2` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
}

func TestEvaluateStringEscape(t *testing.T) {
	context := NewContext(nil)
	if actual, err := context.EvaluateString(`foo \={bar}`); err != nil || actual != `foo ={bar}` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo \\={"bar"}`); err != nil || actual != `foo \bar` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo \\\={bar}`); err != nil || actual != `foo \={bar}` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo \\\\={"bar"}`); err != nil || actual != `foo \\bar` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo \#{bar}`); err != nil || actual != `foo #{bar}` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo \\#{"bar"}`); err != nil || actual != `foo \bar` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo \\\#{bar}`); err != nil || actual != `foo \#{bar}` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo \\\\#{"bar"}`); err != nil || actual != `foo \\bar` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
}

func TestEvaluateStringWithProperties(t *testing.T) {
	properties := map[string]interface{}{
		"FOO": "foo",
		"BAR": "bar",
	}
	build := &Build{
		Dir:        "dir",
		Here:       "here",
		Properties: properties,
	}
	context := NewContext(build)
	context.InitProperties()
	if actual, err := context.EvaluateString(`foo`); err != nil || actual != `foo` {
		t.Errorf("TestEvaluateStringWithProperties failure")
	}
	if actual, err := context.EvaluateString(`={FOO} bar`); err != nil || actual != `foo bar` {
		t.Errorf("TestEvaluateStringWithProperties failure")
	}
	if actual, err := context.EvaluateString(`={FOO} ={BAR}`); err != nil || actual != `foo bar` {
		t.Errorf("TestEvaluateStringWithProperties failure")
	}
	if _, err := context.EvaluateString(`={XXX}`); err == nil || err.Error() != `Undefined symbol 'XXX' (at line 1, column 1)` {
		t.Errorf("TestEvaluateStringWithProperties failure")
	}
	if actual, err := context.EvaluateString(`={if true {"true"\} else {"false"\}}`); err != nil || actual != "true" {
		t.Errorf("TestEvaluateStringWithProperties failure: '%v' - '%s'", err, actual)
	}
}

func TestEvaluateSliceWithProperties(t *testing.T) {
	properties := map[string]interface{}{
		"FOO": "foo",
		"BAR": "bar",
	}
	build := &Build{
		Dir:        "dir",
		Here:       "here",
		Properties: properties,
	}
	context := NewContext(build)
	context.InitProperties()
	actual, err := context.EvaluateObject([]string{`={FOO} BAR`, `FOO ={BAR}`})
	if err != nil {
		t.Fail()
	}
	if reflect.TypeOf(actual) != reflect.SliceOf(reflect.TypeOf("")) {
		t.Fail()
	}
	value := reflect.ValueOf(actual)
	if value.Len() != 2 {
		t.Fail()
	}
	if value.Index(0) == reflect.ValueOf(`foo BAR`) {
		t.Fail()
	}
	if value.Index(1) == reflect.ValueOf(`FOO bar`) {
		t.Fail()
	}
}

func TestEvaluateMapWithProperties(t *testing.T) {
	properties := map[string]interface{}{
		"FOO": "foo",
		"BAR": "bar",
	}
	build := &Build{
		Dir:        "dir",
		Here:       "here",
		Properties: properties,
	}
	context := NewContext(build)
	context.InitProperties()
	actual, err := context.EvaluateObject(map[string]string{"={FOO}": "BAR", "FOO": "={BAR}"})
	if err != nil {
		t.Fail()
	}
	if reflect.TypeOf(actual) != reflect.TypeOf(make(map[string]string)) {
		t.Fail()
	}
	value := reflect.ValueOf(actual)
	if value.Len() != 2 {
		t.Fail()
	}
	if value.MapIndex(reflect.ValueOf("foo")) == reflect.ValueOf("BAR") {
		t.Fail()
	}
	if value.MapIndex(reflect.ValueOf("FOO")) == reflect.ValueOf("bar") {
		t.Fail()
	}
}

func TestGetSetProperty(t *testing.T) {
	context := NewContext(nil)
	context.SetProperty("foo", "bar")
	if p, err := context.GetProperty("foo"); p != "bar" || err != nil {
		t.Fail()
	}
}

func TestEvaluateExpression(t *testing.T) {
	context := NewContext(nil)
	_, err := context.EvaluateExpression(`foo = "BAR"`)
	if err != nil {
		t.Fail()
	}
	r, err := context.GetProperty("foo")
	if err != nil || r != "BAR" {
		t.Fail()
	}
	if r, err = context.EvaluateExpression(`1+2`); err != nil || r != int64(3) {
		t.Fail()
	}
}

func TestEvaluateEnvironmentSimple(t *testing.T) {
	build := &Build{
		Dir: "dir",
		Environment: map[string]string{
			"FOO": "BAR",
		},
	}
	context := NewContext(build)
	env, err := context.EvaluateEnvironment()
	if err != nil {
		t.Errorf("Error getting environment: %v", err)
	}
	for _, line := range env {
		if line == "FOO=BAR" {
			return
		}
	}
	t.Error("Env FOO=BAR not found")
}

func TestEvaluateEnvironmentComplex(t *testing.T) {
	build := &Build{
		Dir: "dir",
		Environment: map[string]string{
			"FOO": "BAR:${HOME}",
		},
	}
	context := NewContext(build)
	env, err := context.EvaluateEnvironment()
	if err != nil {
		t.Errorf("Error getting environment: %v", err)
	}
	var foo string
	for _, line := range env {
		if strings.HasPrefix(line, "HOME=") {
			foo = "FOO=BAR:" + line[5:]
		}
	}
	if foo == "" {
		return
	}
	for _, line := range env {
		if line == foo {
			return
		}
	}
	t.Error("Environment variable FOO not set correctly")
}
<commit_msg>Fixed test for new error message<commit_after>package build

import (
	"reflect"
	"strings"
	"testing"
)

func TestEvaluateString(t *testing.T) {
	context := NewContext(nil)
	if actual, err := context.EvaluateString(`foo`); err != nil || actual != `foo` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo ={`); err != nil || actual != `foo ={` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo #{"bar"}`); err != nil || actual != `foo bar` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo ={"bar"}`); err != nil || actual != `foo bar` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo ={1+1}`); err != nil || actual != `foo 2` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
}

func TestEvaluateStringEscape(t *testing.T) {
	context := NewContext(nil)
	if actual, err := context.EvaluateString(`foo \={bar}`); err != nil || actual != `foo ={bar}` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo \\={"bar"}`); err != nil || actual != `foo \bar` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo \\\={bar}`); err != nil || actual != `foo \={bar}` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo \\\\={"bar"}`); err != nil || actual != `foo \\bar` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo \#{bar}`); err != nil || actual != `foo #{bar}` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo \\#{"bar"}`); err != nil || actual != `foo \bar` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo \\\#{bar}`); err != nil || actual != `foo \#{bar}` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
	if actual, err := context.EvaluateString(`foo \\\\#{"bar"}`); err != nil || actual != `foo \\bar` {
		t.Errorf("TestEvaluateString failure: '%s'", actual)
	}
}

func TestEvaluateStringWithProperties(t *testing.T) {
	properties := map[string]interface{}{
		"FOO": "foo",
		"BAR": "bar",
	}
	build := &Build{
		Dir:        "dir",
		Here:       "here",
		Properties: properties,
	}
	context := NewContext(build)
	context.InitProperties()
	if actual, err := context.EvaluateString(`foo`); err != nil || actual != `foo` {
		t.Errorf("TestEvaluateStringWithProperties failure")
	}
	if actual, err := context.EvaluateString(`={FOO} bar`); err != nil || actual != `foo bar` {
		t.Errorf("TestEvaluateStringWithProperties failure")
	}
	if actual, err := context.EvaluateString(`={FOO} ={BAR}`); err != nil || actual != `foo bar` {
		t.Errorf("TestEvaluateStringWithProperties failure")
	}
	if _, err := context.EvaluateString(`={XXX}`); err == nil || err.Error() != `undefined symbol 'XXX' (at line 1, column 1)` {
		t.Errorf("TestEvaluateStringWithProperties failure: %v", err)
	}
	if actual, err := context.EvaluateString(`={if true {"true"\} else {"false"\}}`); err != nil || actual != "true" {
		t.Errorf("TestEvaluateStringWithProperties failure: '%v' - '%s'", err, actual)
	}
}

func TestEvaluateSliceWithProperties(t *testing.T) {
	properties := map[string]interface{}{
		"FOO": "foo",
		"BAR": "bar",
	}
	build := &Build{
		Dir:        "dir",
		Here:       "here",
		Properties: properties,
	}
	context := NewContext(build)
	context.InitProperties()
	actual, err := context.EvaluateObject([]string{`={FOO} BAR`, `FOO ={BAR}`})
	if err != nil {
		t.Fail()
	}
	if reflect.TypeOf(actual) != reflect.SliceOf(reflect.TypeOf("")) {
		t.Fail()
	}
	value := reflect.ValueOf(actual)
	if value.Len() != 2 {
		t.Fail()
	}
	if value.Index(0) == reflect.ValueOf(`foo BAR`) {
		t.Fail()
	}
	if value.Index(1) == reflect.ValueOf(`FOO bar`) {
		t.Fail()
	}
}

func TestEvaluateMapWithProperties(t *testing.T) {
	properties := map[string]interface{}{
		"FOO": "foo",
		"BAR": "bar",
	}
	build := &Build{
		Dir:        "dir",
		Here:       "here",
		Properties: properties,
	}
	context := NewContext(build)
	context.InitProperties()
	actual, err := context.EvaluateObject(map[string]string{"={FOO}": "BAR", "FOO": "={BAR}"})
	if err != nil {
		t.Fail()
	}
	if reflect.TypeOf(actual) != reflect.TypeOf(make(map[string]string)) {
		t.Fail()
	}
	value := reflect.ValueOf(actual)
	if value.Len() != 2 {
		t.Fail()
	}
	if value.MapIndex(reflect.ValueOf("foo")) == reflect.ValueOf("BAR") {
		t.Fail()
	}
	if value.MapIndex(reflect.ValueOf("FOO")) == reflect.ValueOf("bar") {
		t.Fail()
	}
}

func TestGetSetProperty(t *testing.T) {
	context := NewContext(nil)
	context.SetProperty("foo", "bar")
	if p, err := context.GetProperty("foo"); p != "bar" || err != nil {
		t.Fail()
	}
}

func TestEvaluateExpression(t *testing.T) {
	context := NewContext(nil)
	_, err := context.EvaluateExpression(`foo = "BAR"`)
	if err != nil {
		t.Fail()
	}
	r, err := context.GetProperty("foo")
	if err != nil || r != "BAR" {
		t.Fail()
	}
	if r, err = context.EvaluateExpression(`1+2`); err != nil || r != int64(3) {
		t.Fail()
	}
}

func TestEvaluateEnvironmentSimple(t *testing.T) {
	build := &Build{
		Dir: "dir",
		Environment: map[string]string{
			"FOO": "BAR",
		},
	}
	context := NewContext(build)
	env, err := context.EvaluateEnvironment()
	if err != nil {
		t.Errorf("Error getting environment: %v", err)
	}
	for _, line := range env {
		if line == "FOO=BAR" {
			return
		}
	}
	t.Error("Env FOO=BAR not found")
}

func TestEvaluateEnvironmentComplex(t *testing.T) {
	build := &Build{
		Dir: "dir",
		Environment: map[string]string{
			"FOO": "BAR:${HOME}",
		},
	}
	context := NewContext(build)
	env, err := context.EvaluateEnvironment()
	if err != nil {
		t.Errorf("Error getting environment: %v", err)
	}
	var foo string
	for _, line := range env {
		if strings.HasPrefix(line, "HOME=") {
			foo = "FOO=BAR:" + line[5:]
		}
	}
	if foo == "" {
		return
	}
	for _, line := range env {
		if line == foo {
			return
		}
	}
	t.Error("Environment variable FOO not set correctly")
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"femebe"
	"io"
	"log"
	"logplexc"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"time"
)

const (
	KB uint32 = 1024
	MB        = 1048576
)

func logWorker(rwc io.ReadWriteCloser, cfg logplexc.Config) {
	var m femebe.Message
	var err error
	stream := femebe.NewServerMessageStream("", rwc)

	// Exit in an orderly manner if (and only if) exitFn() is
	// called.
	defer func() {
		rwc.Close()

		if r := recover(); r != nil && r != &m {
			panic(r)
		}
	}()

	// Convention to exit the function; useful in closures that
	// may need to jump block scopes.
	exitFn := func(args ...interface{}) {
		if len(args) == 1 {
			log.Printf("Disconnect client: %v", args[0])
		} else if len(args) > 1 {
			if s, ok := args[0].(string); ok {
				log.Printf(s, args[1:])
			} else {
				// Not an intended use case, but do
				// one's best to print something.
				log.Printf("Got a malformed exit: %v", args)
			}
		}

		panic(&m)
	}

	// Function to get the next message and fill m, or immediately
	// exit the parent function after logging the offending error.
	mustFillNext := func() {
		err = stream.Next(&m)
		if err != nil {
			log.Printf("could not read next message: %v", err)
			exitFn(err)
		}
	}

	// Read the version message and bomb out if it looks funny in
	// any way.  This is also the area that will need amendment as
	// new versions of the protocol and new versions of Postgres
	// (which may change what error fields are available) are
	// supported.
	mustCheckVersion := func() {
		if m.MsgType() != 'V' {
			log.Printf("expected version ('V') message, "+
				"but received %c", m.MsgType())
			exitFn(nil)
		}

		// hard-coded lengh limit, but it's very generous
		if m.Size() > 10*KB {
			log.Printf("oversized message string, msg size is %d",
				m.Size())
		}

		s, err := femebe.ReadCString(m.Payload())
		if err != nil {
			log.Printf("couldn't read version string: %v", err)
			exitFn(nil)
		}

		if !strings.HasPrefix(s, "PG9.2.") ||
			!strings.HasSuffix(s, "/1") {
			exitFn("protocol version not supported: %s", s)
		}
	}

	// Read the remote system identifier string
	mustReadIdentifier := func() string {
		if m.MsgType() != 'I' {
			log.Printf("expected identification ('I') message, "+
				"but received %c", m.MsgType())
			exitFn()
		}

		// hard-coded lengh limit, but it's very generous
		if m.Size() > 10*KB {
			log.Printf("oversized message string, msg size is %d",
				m.Size())
		}

		s, err := femebe.ReadCString(m.Payload())
		if err != nil {
			log.Printf("couldn't read identification string: %v",
				err)
			exitFn()
		}

		return s
	}

	// Protocol start-up; packets that are only received once.
	mustFillNext()
	mustCheckVersion()
	mustFillNext()
	ident := mustReadIdentifier()
	log.Printf("client connects with identifier %q", ident)

	cfg.Token = ident
	client, err := logplexc.NewClient(&cfg)
	if err != nil {
		exitFn(err)
	}

	for {
		mustFillNext()

		// Refuse to handle any log message above an arbitrary
		// size.  Furthermore, exit the worker, closing the
		// connection, so that the client doesn't even bother
		// to wait for this process to drain the oversized
		// item and anything following it; these will be
		// dropped.  It's on the client to gracefully handle
		// the error and re-connect after this happens.
		if m.Size() > 1*MB {
			exitFn("client %q sent oversized log record, "+
				"disconnecting", ident)
		}

		payload, err := m.Force()
		if err != nil {
			exitFn("could not retreive payload of message: %v",
				err)
		}

		var lr logRecord
		parseLogRecord(&lr, payload, exitFn)

		// Buffer to format the complete log message in.
		msgFmtBuf := bytes.Buffer{}

		// Helps with formatting a series of nullable strings.
		catOptionalField := func(prefix string, maybePresent *string) {
			if maybePresent != nil {
				if prefix != "" {
					msgFmtBuf.WriteString(prefix)
					msgFmtBuf.WriteString(": ")
				}

				msgFmtBuf.WriteString(*maybePresent)
				msgFmtBuf.WriteByte('\n')
			}
		}

		catOptionalField("", lr.ErrMessage)
		catOptionalField("Detail", lr.ErrDetail)
		catOptionalField("Hint", lr.ErrHint)
		catOptionalField("Query", lr.UserQuery)

		client.BufferMessage(time.Now(), "postgres-"+lr.SessionId,
			msgFmtBuf.Bytes())

		// Post the buffered message
		resp, err := client.PostMessages()
		if err != nil {
			exitFn(err)
		}

		// Can be more carefully checked and reported on to
		// indicate configuration or Logplex issues.
		resp.Body.Close()
	}
}

func main() {
	// Input checking
	if len(os.Args) != 2 {
		log.Printf("Usage: pg_logplexcollector LISTENADDR\n")
		os.Exit(1)
	}

	// Set up log prefix for all future system diagnostic
	// messages.
	log.SetPrefix("pg_logplexcollector: ")

	// Signal handling: print dying gasp and and exit
	sigch := make(chan os.Signal)
	signal.Notify(sigch, os.Interrupt, os.Kill)
	go func() {
		for sig := range sigch {
			log.Printf("got signal %v", sig)
			if sig == os.Kill {
				os.Exit(2)
			} else if sig == os.Interrupt {
				os.Exit(0)
			}
		}
	}()

	logplexUrl, err := url.Parse(os.Getenv("LOGPLEX_URL"))
	if err != nil {
		log.Fatalf("Could not parse logplex endpoint %q: %v",
			os.Getenv("LOGPLEX_URL"), err)
	}

	templateConfig := logplexc.Config{
		Logplex:    *logplexUrl,
		Token:      "",
		HttpClient: *http.DefaultClient,
		Transport: http.Transport{
			TLSClientConfig: &tls.Config{
				InsecureSkipVerify: true,
			},
		},
	}

	// Begin listening
	l, err := net.Listen("unix", os.Args[1])
	if err != nil {
		log.Fatalf(
			"exiting, cannot listen to %q: %v",
			os.Args[1], err)
	}

	for {
		conn, err := l.Accept()
		if err != nil {
			log.Printf("accept error: %v", err)
		}

		go logWorker(conn, templateConfig)
	}
}
<commit_msg>Change log prefix slightly<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"femebe"
	"io"
	"log"
	"logplexc"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"time"
)

const (
	KB uint32 = 1024
	MB        = 1048576
)

func logWorker(rwc io.ReadWriteCloser, cfg logplexc.Config) {
	var m femebe.Message
	var err error
	stream := femebe.NewServerMessageStream("", rwc)

	// Exit in an orderly manner if (and only if) exitFn() is
	// called.
	defer func() {
		rwc.Close()

		if r := recover(); r != nil && r != &m {
			panic(r)
		}
	}()

	// Convention to exit the function; useful in closures that
	// may need to jump block scopes.
	exitFn := func(args ...interface{}) {
		if len(args) == 1 {
			log.Printf("Disconnect client: %v", args[0])
		} else if len(args) > 1 {
			if s, ok := args[0].(string); ok {
				log.Printf(s, args[1:])
			} else {
				// Not an intended use case, but do
				// one's best to print something.
				log.Printf("Got a malformed exit: %v", args)
			}
		}

		panic(&m)
	}

	// Function to get the next message and fill m, or immediately
	// exit the parent function after logging the offending error.
	mustFillNext := func() {
		err = stream.Next(&m)
		if err != nil {
			log.Printf("could not read next message: %v", err)
			exitFn(err)
		}
	}

	// Read the version message and bomb out if it looks funny in
	// any way.  This is also the area that will need amendment as
	// new versions of the protocol and new versions of Postgres
	// (which may change what error fields are available) are
	// supported.
	mustCheckVersion := func() {
		if m.MsgType() != 'V' {
			log.Printf("expected version ('V') message, "+
				"but received %c", m.MsgType())
			exitFn(nil)
		}

		// hard-coded lengh limit, but it's very generous
		if m.Size() > 10*KB {
			log.Printf("oversized message string, msg size is %d",
				m.Size())
		}

		s, err := femebe.ReadCString(m.Payload())
		if err != nil {
			log.Printf("couldn't read version string: %v", err)
			exitFn(nil)
		}

		if !strings.HasPrefix(s, "PG9.2.") ||
			!strings.HasSuffix(s, "/1") {
			exitFn("protocol version not supported: %s", s)
		}
	}

	// Read the remote system identifier string
	mustReadIdentifier := func() string {
		if m.MsgType() != 'I' {
			log.Printf("expected identification ('I') message, "+
				"but received %c", m.MsgType())
			exitFn()
		}

		// hard-coded lengh limit, but it's very generous
		if m.Size() > 10*KB {
			log.Printf("oversized message string, msg size is %d",
				m.Size())
		}

		s, err := femebe.ReadCString(m.Payload())
		if err != nil {
			log.Printf("couldn't read identification string: %v",
				err)
			exitFn()
		}

		return s
	}

	// Protocol start-up; packets that are only received once.
	mustFillNext()
	mustCheckVersion()
	mustFillNext()
	ident := mustReadIdentifier()
	log.Printf("client connects with identifier %q", ident)

	cfg.Token = ident
	client, err := logplexc.NewClient(&cfg)
	if err != nil {
		exitFn(err)
	}

	for {
		mustFillNext()

		// Refuse to handle any log message above an arbitrary
		// size.  Furthermore, exit the worker, closing the
		// connection, so that the client doesn't even bother
		// to wait for this process to drain the oversized
		// item and anything following it; these will be
		// dropped.  It's on the client to gracefully handle
		// the error and re-connect after this happens.
		if m.Size() > 1*MB {
			exitFn("client %q sent oversized log record, "+
				"disconnecting", ident)
		}

		payload, err := m.Force()
		if err != nil {
			exitFn("could not retreive payload of message: %v",
				err)
		}

		var lr logRecord
		parseLogRecord(&lr, payload, exitFn)

		// Buffer to format the complete log message in.
		msgFmtBuf := bytes.Buffer{}

		// Helps with formatting a series of nullable strings.
		catOptionalField := func(prefix string, maybePresent *string) {
			if maybePresent != nil {
				if prefix != "" {
					msgFmtBuf.WriteString(prefix)
					msgFmtBuf.WriteString(": ")
				}

				msgFmtBuf.WriteString(*maybePresent)
				msgFmtBuf.WriteByte('\n')
			}
		}

		catOptionalField("", lr.ErrMessage)
		catOptionalField("Detail", lr.ErrDetail)
		catOptionalField("Hint", lr.ErrHint)
		catOptionalField("Query", lr.UserQuery)

		client.BufferMessage(time.Now(), "postgres-"+lr.SessionId,
			msgFmtBuf.Bytes())

		// Post the buffered message
		resp, err := client.PostMessages()
		if err != nil {
			exitFn(err)
		}

		// Can be more carefully checked and reported on to
		// indicate configuration or Logplex issues.
		resp.Body.Close()
	}
}

func main() {
	// Input checking
	if len(os.Args) != 2 {
		log.Printf("Usage: pg_logplexcollector LISTENADDR\n")
		os.Exit(1)
	}

	// Set up log prefix for all future system diagnostic
	// messages.
	log.SetPrefix("pg_logplexcollector ")

	// Signal handling: print dying gasp and and exit
	sigch := make(chan os.Signal)
	signal.Notify(sigch, os.Interrupt, os.Kill)
	go func() {
		for sig := range sigch {
			log.Printf("got signal %v", sig)
			if sig == os.Kill {
				os.Exit(2)
			} else if sig == os.Interrupt {
				os.Exit(0)
			}
		}
	}()

	logplexUrl, err := url.Parse(os.Getenv("LOGPLEX_URL"))
	if err != nil {
		log.Fatalf("Could not parse logplex endpoint %q: %v",
			os.Getenv("LOGPLEX_URL"), err)
	}

	templateConfig := logplexc.Config{
		Logplex:    *logplexUrl,
		Token:      "",
		HttpClient: *http.DefaultClient,
		Transport: http.Transport{
			TLSClientConfig: &tls.Config{
				InsecureSkipVerify: true,
			},
		},
	}

	// Begin listening
	l, err := net.Listen("unix", os.Args[1])
	if err != nil {
		log.Fatalf(
			"exiting, cannot listen to %q: %v",
			os.Args[1], err)
	}

	for {
		conn, err := l.Accept()
		if err != nil {
			log.Printf("accept error: %v", err)
		}

		go logWorker(conn, templateConfig)
	}
}
<|endoftext|>


<commit_before>// Copyright 2010 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.

// +build darwin dragonfly freebsd linux nacl netbsd openbsd solaris windows

package net

import (
	"syscall"
	"time"
)

// BUG(mikio): On every POSIX platform, reads from the "ip4" network
// using the ReadFrom or ReadFromIP method might not return a complete
// IPv4 packet, including its header, even if there is space
// available. This can occur even in cases where Read or ReadMsgIP
// could return a complete packet. For this reason, it is recommended
// that you do not uses these methods if it is important to receive a
// full packet.
//
// The Go 1 compatibility guidelines make it impossible for us to
// change the behavior of these methods; use Read or ReadMsgIP
// instead.

func sockaddrToIP(sa syscall.Sockaddr) Addr {
	switch sa := sa.(type) {
	case *syscall.SockaddrInet4:
		return &IPAddr{IP: sa.Addr[0:]}
	case *syscall.SockaddrInet6:
		return &IPAddr{IP: sa.Addr[0:], Zone: zoneToString(int(sa.ZoneId))}
	}
	return nil
}

func (a *IPAddr) family() int {
	if a == nil || len(a.IP) <= IPv4len {
		return syscall.AF_INET
	}
	if a.IP.To4() != nil {
		return syscall.AF_INET
	}
	return syscall.AF_INET6
}

func (a *IPAddr) isWildcard() bool {
	if a == nil || a.IP == nil {
		return true
	}
	return a.IP.IsUnspecified()
}

func (a *IPAddr) sockaddr(family int) (syscall.Sockaddr, error) {
	if a == nil {
		return nil, nil
	}
	return ipToSockaddr(family, a.IP, 0, a.Zone)
}

// IPConn is the implementation of the Conn and PacketConn interfaces
// for IP network connections.
type IPConn struct {
	conn
}

func newIPConn(fd *netFD) *IPConn { return &IPConn{conn{fd}} }

// ReadFromIP reads an IP packet from c, copying the payload into b.
// It returns the number of bytes copied into b and the return address
// that was on the packet.
//
// ReadFromIP can be made to time out and return an error with
// Timeout() == true after a fixed time limit; see SetDeadline and
// SetReadDeadline.
func (c *IPConn) ReadFromIP(b []byte) (int, *IPAddr, error) {
	if !c.ok() {
		return 0, nil, syscall.EINVAL
	}
	// TODO(cw,rsc): consider using readv if we know the family
	// type to avoid the header trim/copy
	var addr *IPAddr
	n, sa, err := c.fd.readFrom(b)
	switch sa := sa.(type) {
	case *syscall.SockaddrInet4:
		addr = &IPAddr{IP: sa.Addr[0:]}
		if len(b) >= IPv4len { // discard ipv4 header
			hsize := (int(b[0]) & 0xf) * 4
			copy(b, b[hsize:])
			n -= hsize
		}
	case *syscall.SockaddrInet6:
		addr = &IPAddr{IP: sa.Addr[0:], Zone: zoneToString(int(sa.ZoneId))}
	}
	return n, addr, err
}

// ReadFrom implements the PacketConn ReadFrom method.
func (c *IPConn) ReadFrom(b []byte) (int, Addr, error) {
	if !c.ok() {
		return 0, nil, syscall.EINVAL
	}
	n, addr, err := c.ReadFromIP(b)
	return n, addr.toAddr(), err
}

// ReadMsgIP reads a packet from c, copying the payload into b and the
// associated out-of-band data into oob.  It returns the number of
// bytes copied into b, the number of bytes copied into oob, the flags
// that were set on the packet and the source address of the packet.
func (c *IPConn) ReadMsgIP(b, oob []byte) (n, oobn, flags int, addr *IPAddr, err error) {
	if !c.ok() {
		return 0, 0, 0, nil, syscall.EINVAL
	}
	var sa syscall.Sockaddr
	n, oobn, flags, sa, err = c.fd.readMsg(b, oob)
	switch sa := sa.(type) {
	case *syscall.SockaddrInet4:
		addr = &IPAddr{IP: sa.Addr[0:]}
	case *syscall.SockaddrInet6:
		addr = &IPAddr{IP: sa.Addr[0:], Zone: zoneToString(int(sa.ZoneId))}
	}
	return
}

// WriteToIP writes an IP packet to addr via c, copying the payload
// from b.
//
// WriteToIP can be made to time out and return an error with
// Timeout() == true after a fixed time limit; see SetDeadline and
// SetWriteDeadline.  On packet-oriented connections, write timeouts
// are rare.
func (c *IPConn) WriteToIP(b []byte, addr *IPAddr) (int, error) {
	if !c.ok() {
		return 0, syscall.EINVAL
	}
	if addr == nil {
		return 0, &OpError{Op: "write", Net: c.fd.net, Addr: nil, Err: errMissingAddress}
	}
	sa, err := addr.sockaddr(c.fd.family)
	if err != nil {
		return 0, &OpError{"write", c.fd.net, addr, err}
	}
	return c.fd.writeTo(b, sa)
}

// WriteTo implements the PacketConn WriteTo method.
func (c *IPConn) WriteTo(b []byte, addr Addr) (int, error) {
	if !c.ok() {
		return 0, syscall.EINVAL
	}
	a, ok := addr.(*IPAddr)
	if !ok {
		return 0, &OpError{"write", c.fd.net, addr, syscall.EINVAL}
	}
	return c.WriteToIP(b, a)
}

// WriteMsgIP writes a packet to addr via c, copying the payload from
// b and the associated out-of-band data from oob.  It returns the
// number of payload and out-of-band bytes written.
func (c *IPConn) WriteMsgIP(b, oob []byte, addr *IPAddr) (n, oobn int, err error) {
	if !c.ok() {
		return 0, 0, syscall.EINVAL
	}
	if addr == nil {
		return 0, 0, &OpError{Op: "write", Net: c.fd.net, Addr: nil, Err: errMissingAddress}
	}
	sa, err := addr.sockaddr(c.fd.family)
	if err != nil {
		return 0, 0, &OpError{"write", c.fd.net, addr, err}
	}
	return c.fd.writeMsg(b, oob, sa)
}

// DialIP connects to the remote address raddr on the network protocol
// netProto, which must be "ip", "ip4", or "ip6" followed by a colon
// and a protocol number or name.
func DialIP(netProto string, laddr, raddr *IPAddr) (*IPConn, error) {
	return dialIP(netProto, laddr, raddr, noDeadline)
}

func dialIP(netProto string, laddr, raddr *IPAddr, deadline time.Time) (*IPConn, error) {
	net, proto, err := parseNetwork(netProto)
	if err != nil {
		return nil, &OpError{Op: "dial", Net: netProto, Addr: raddr, Err: err}
	}
	switch net {
	case "ip", "ip4", "ip6":
	default:
		return nil, &OpError{Op: "dial", Net: netProto, Addr: raddr, Err: UnknownNetworkError(netProto)}
	}
	if raddr == nil {
		return nil, &OpError{Op: "dial", Net: netProto, Addr: nil, Err: errMissingAddress}
	}
	fd, err := internetSocket(net, laddr, raddr, deadline, syscall.SOCK_RAW, proto, "dial", sockaddrToIP)
	if err != nil {
		return nil, &OpError{Op: "dial", Net: netProto, Addr: raddr, Err: err}
	}
	return newIPConn(fd), nil
}

// ListenIP listens for incoming IP packets addressed to the local
// address laddr.  The returned connection's ReadFrom and WriteTo
// methods can be used to receive and send IP packets with per-packet
// addressing.
func ListenIP(netProto string, laddr *IPAddr) (*IPConn, error) {
	net, proto, err := parseNetwork(netProto)
	if err != nil {
		return nil, &OpError{Op: "dial", Net: netProto, Addr: laddr, Err: err}
	}
	switch net {
	case "ip", "ip4", "ip6":
	default:
		return nil, &OpError{Op: "listen", Net: netProto, Addr: laddr, Err: UnknownNetworkError(netProto)}
	}
	fd, err := internetSocket(net, laddr, nil, noDeadline, syscall.SOCK_RAW, proto, "listen", sockaddrToIP)
	if err != nil {
		return nil, &OpError{Op: "listen", Net: netProto, Addr: laddr, Err: err}
	}
	return newIPConn(fd), nil
}
<commit_msg>net: make WriteTo, WriteToIP and WriteMsgIP fail when IPConn is already connected<commit_after>// Copyright 2010 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.

// +build darwin dragonfly freebsd linux nacl netbsd openbsd solaris windows

package net

import (
	"syscall"
	"time"
)

// BUG(mikio): On every POSIX platform, reads from the "ip4" network
// using the ReadFrom or ReadFromIP method might not return a complete
// IPv4 packet, including its header, even if there is space
// available. This can occur even in cases where Read or ReadMsgIP
// could return a complete packet. For this reason, it is recommended
// that you do not uses these methods if it is important to receive a
// full packet.
//
// The Go 1 compatibility guidelines make it impossible for us to
// change the behavior of these methods; use Read or ReadMsgIP
// instead.

func sockaddrToIP(sa syscall.Sockaddr) Addr {
	switch sa := sa.(type) {
	case *syscall.SockaddrInet4:
		return &IPAddr{IP: sa.Addr[0:]}
	case *syscall.SockaddrInet6:
		return &IPAddr{IP: sa.Addr[0:], Zone: zoneToString(int(sa.ZoneId))}
	}
	return nil
}

func (a *IPAddr) family() int {
	if a == nil || len(a.IP) <= IPv4len {
		return syscall.AF_INET
	}
	if a.IP.To4() != nil {
		return syscall.AF_INET
	}
	return syscall.AF_INET6
}

func (a *IPAddr) isWildcard() bool {
	if a == nil || a.IP == nil {
		return true
	}
	return a.IP.IsUnspecified()
}

func (a *IPAddr) sockaddr(family int) (syscall.Sockaddr, error) {
	if a == nil {
		return nil, nil
	}
	return ipToSockaddr(family, a.IP, 0, a.Zone)
}

// IPConn is the implementation of the Conn and PacketConn interfaces
// for IP network connections.
type IPConn struct {
	conn
}

func newIPConn(fd *netFD) *IPConn { return &IPConn{conn{fd}} }

// ReadFromIP reads an IP packet from c, copying the payload into b.
// It returns the number of bytes copied into b and the return address
// that was on the packet.
//
// ReadFromIP can be made to time out and return an error with
// Timeout() == true after a fixed time limit; see SetDeadline and
// SetReadDeadline.
func (c *IPConn) ReadFromIP(b []byte) (int, *IPAddr, error) {
	if !c.ok() {
		return 0, nil, syscall.EINVAL
	}
	// TODO(cw,rsc): consider using readv if we know the family
	// type to avoid the header trim/copy
	var addr *IPAddr
	n, sa, err := c.fd.readFrom(b)
	switch sa := sa.(type) {
	case *syscall.SockaddrInet4:
		addr = &IPAddr{IP: sa.Addr[0:]}
		if len(b) >= IPv4len { // discard ipv4 header
			hsize := (int(b[0]) & 0xf) * 4
			copy(b, b[hsize:])
			n -= hsize
		}
	case *syscall.SockaddrInet6:
		addr = &IPAddr{IP: sa.Addr[0:], Zone: zoneToString(int(sa.ZoneId))}
	}
	return n, addr, err
}

// ReadFrom implements the PacketConn ReadFrom method.
func (c *IPConn) ReadFrom(b []byte) (int, Addr, error) {
	if !c.ok() {
		return 0, nil, syscall.EINVAL
	}
	n, addr, err := c.ReadFromIP(b)
	return n, addr.toAddr(), err
}

// ReadMsgIP reads a packet from c, copying the payload into b and the
// associated out-of-band data into oob.  It returns the number of
// bytes copied into b, the number of bytes copied into oob, the flags
// that were set on the packet and the source address of the packet.
func (c *IPConn) ReadMsgIP(b, oob []byte) (n, oobn, flags int, addr *IPAddr, err error) {
	if !c.ok() {
		return 0, 0, 0, nil, syscall.EINVAL
	}
	var sa syscall.Sockaddr
	n, oobn, flags, sa, err = c.fd.readMsg(b, oob)
	switch sa := sa.(type) {
	case *syscall.SockaddrInet4:
		addr = &IPAddr{IP: sa.Addr[0:]}
	case *syscall.SockaddrInet6:
		addr = &IPAddr{IP: sa.Addr[0:], Zone: zoneToString(int(sa.ZoneId))}
	}
	return
}

// WriteToIP writes an IP packet to addr via c, copying the payload
// from b.
//
// WriteToIP can be made to time out and return an error with
// Timeout() == true after a fixed time limit; see SetDeadline and
// SetWriteDeadline.  On packet-oriented connections, write timeouts
// are rare.
func (c *IPConn) WriteToIP(b []byte, addr *IPAddr) (int, error) {
	if !c.ok() {
		return 0, syscall.EINVAL
	}
	if c.fd.isConnected {
		return 0, &OpError{Op: "write", Net: c.fd.net, Addr: addr, Err: ErrWriteToConnected}
	}
	if addr == nil {
		return 0, &OpError{Op: "write", Net: c.fd.net, Addr: nil, Err: errMissingAddress}
	}
	sa, err := addr.sockaddr(c.fd.family)
	if err != nil {
		return 0, &OpError{"write", c.fd.net, addr, err}
	}
	return c.fd.writeTo(b, sa)
}

// WriteTo implements the PacketConn WriteTo method.
func (c *IPConn) WriteTo(b []byte, addr Addr) (int, error) {
	if !c.ok() {
		return 0, syscall.EINVAL
	}
	a, ok := addr.(*IPAddr)
	if !ok {
		return 0, &OpError{"write", c.fd.net, addr, syscall.EINVAL}
	}
	return c.WriteToIP(b, a)
}

// WriteMsgIP writes a packet to addr via c, copying the payload from
// b and the associated out-of-band data from oob.  It returns the
// number of payload and out-of-band bytes written.
func (c *IPConn) WriteMsgIP(b, oob []byte, addr *IPAddr) (n, oobn int, err error) {
	if !c.ok() {
		return 0, 0, syscall.EINVAL
	}
	if c.fd.isConnected {
		return 0, 0, &OpError{Op: "write", Net: c.fd.net, Addr: addr, Err: ErrWriteToConnected}
	}
	if addr == nil {
		return 0, 0, &OpError{Op: "write", Net: c.fd.net, Addr: nil, Err: errMissingAddress}
	}
	sa, err := addr.sockaddr(c.fd.family)
	if err != nil {
		return 0, 0, &OpError{"write", c.fd.net, addr, err}
	}
	return c.fd.writeMsg(b, oob, sa)
}

// DialIP connects to the remote address raddr on the network protocol
// netProto, which must be "ip", "ip4", or "ip6" followed by a colon
// and a protocol number or name.
func DialIP(netProto string, laddr, raddr *IPAddr) (*IPConn, error) {
	return dialIP(netProto, laddr, raddr, noDeadline)
}

func dialIP(netProto string, laddr, raddr *IPAddr, deadline time.Time) (*IPConn, error) {
	net, proto, err := parseNetwork(netProto)
	if err != nil {
		return nil, &OpError{Op: "dial", Net: netProto, Addr: raddr, Err: err}
	}
	switch net {
	case "ip", "ip4", "ip6":
	default:
		return nil, &OpError{Op: "dial", Net: netProto, Addr: raddr, Err: UnknownNetworkError(netProto)}
	}
	if raddr == nil {
		return nil, &OpError{Op: "dial", Net: netProto, Addr: nil, Err: errMissingAddress}
	}
	fd, err := internetSocket(net, laddr, raddr, deadline, syscall.SOCK_RAW, proto, "dial", sockaddrToIP)
	if err != nil {
		return nil, &OpError{Op: "dial", Net: netProto, Addr: raddr, Err: err}
	}
	return newIPConn(fd), nil
}

// ListenIP listens for incoming IP packets addressed to the local
// address laddr.  The returned connection's ReadFrom and WriteTo
// methods can be used to receive and send IP packets with per-packet
// addressing.
func ListenIP(netProto string, laddr *IPAddr) (*IPConn, error) {
	net, proto, err := parseNetwork(netProto)
	if err != nil {
		return nil, &OpError{Op: "dial", Net: netProto, Addr: laddr, Err: err}
	}
	switch net {
	case "ip", "ip4", "ip6":
	default:
		return nil, &OpError{Op: "listen", Net: netProto, Addr: laddr, Err: UnknownNetworkError(netProto)}
	}
	fd, err := internetSocket(net, laddr, nil, noDeadline, syscall.SOCK_RAW, proto, "listen", sockaddrToIP)
	if err != nil {
		return nil, &OpError{Op: "listen", Net: netProto, Addr: laddr, Err: err}
	}
	return newIPConn(fd), nil
}
<|endoftext|>