<commit_before>package http

import (
	"bytes"
	"context"
	"encoding/gob"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"mime"
	"net/http"
	"strings"
)

const (
	// AcceptTypeKey is the context key used to store the value of the HTTP
	// request Accept-Type header. The value may be used by encoders and
	// decoders to implement a content type negotiation algorithm.
	AcceptTypeKey contextKey = iota + 1
	// ContentTypeKey is the context key used to store the value of the HTTP
	// response Content-Type header when explicitly set in the DSL. The value
	// may be used by encoders to set the header appropriately.
	ContentTypeKey
)

type (
	// Decoder provides the actual decoding algorithm used to load HTTP
	// request and response bodies.
	Decoder interface {
		// Decode decodes into v.
		Decode(v interface{}) error
	}

	// Encoder provides the actual encoding algorithm used to write HTTP
	// request and response bodies.
	Encoder interface {
		// Encode encodes v.
		Encode(v interface{}) error
	}

	// EncodingFunc allows a function with appropriate signature to act as a
	// Decoder/Encoder.
	EncodingFunc func(v interface{}) error

	// private type used to define context keys.
	contextKey int
)

// RequestDecoder returns a HTTP request body decoder suitable for the given
// request. The decoder handles the following mime types:
//
//     * application/json using package encoding/json
//     * application/xml using package encoding/xml
//     * application/gob using package encoding/gob
//
// RequestDecoder defaults to the JSON decoder if the request "Content-Type"
// header does not match any of the supported mime type or is missing
// altogether.
func RequestDecoder(r *http.Request) Decoder {
	contentType := r.Header.Get("Content-Type")
	if contentType == "" {
		// default to JSON
		contentType = "application/json"
	} else {
		// sanitize
		if mediaType, _, err := mime.ParseMediaType(contentType); err == nil {
			contentType = mediaType
		}
	}
	switch contentType {
	case "application/json":
		return json.NewDecoder(r.Body)
	case "application/gob":
		return gob.NewDecoder(r.Body)
	case "application/xml":
		return xml.NewDecoder(r.Body)
	default:
		return json.NewDecoder(r.Body)
	}
}

// ResponseEncoder returns a HTTP response encoder leveraging the mime type
// set in the context under the AcceptTypeKey or the ContentTypeKey if any.
// The encoder supports the following mime types:
//
//     * application/json using package encoding/json
//     * application/xml using package encoding/xml
//     * application/gob using package encoding/gob
//     * text/html and text/plain for strings
//
// ResponseEncoder defaults to the JSON encoder if the context AcceptTypeKey or
// ContentTypeKey value does not match any of the supported mime types or is
// missing altogether.
func ResponseEncoder(ctx context.Context, w http.ResponseWriter) Encoder {
	negotiate := func(a string) (Encoder, string) {
		switch a {
		case "", "application/json":
			// default to JSON
			return json.NewEncoder(w), "application/json"
		case "application/xml":
			return xml.NewEncoder(w), "application/xml"
		case "application/gob":
			return gob.NewEncoder(w), "application/gob"
		case "text/html", "text/plain":
			return newTextEncoder(w, a), a
		}
		return nil, ""
	}
	var accept string
	{
		if a := ctx.Value(AcceptTypeKey); a != nil {
			accept = a.(string)
		}
	}
	var ct string
	{
		if a := ctx.Value(ContentTypeKey); a != nil {
			ct = a.(string)
		}
	}
	var (
		enc Encoder
		mt  string
		err error
	)
	{
		if ct != "" {
			// If content type explicitly set in the DSL, infer the response encoder
			// from the content type context key.
			if mt, _, err = mime.ParseMediaType(ct); err == nil {
				switch {
				case ct == "application/json" || strings.HasSuffix(ct, "+json"):
					enc = json.NewEncoder(w)
				case ct == "application/xml" || strings.HasSuffix(ct, "+xml"):
					enc = xml.NewEncoder(w)
				case ct == "application/gob" || strings.HasSuffix(ct, "+gob"):
					enc = gob.NewEncoder(w)
				case ct == "text/html" || ct == "text/plain" ||
					strings.HasSuffix(ct, "+html") || strings.HasSuffix(ct, "+txt"):
					enc = newTextEncoder(w, ct)
				default:
					enc = json.NewEncoder(w)
				}
			}
			SetContentType(w, mt)
			return enc
		}
		// If Accept header exists in the request, infer the response encoder
		// from the header value.
		if enc, mt = negotiate(accept); enc == nil {
			// attempt to normalize
			if mt, _, err = mime.ParseMediaType(accept); err == nil {
				enc, mt = negotiate(mt)
			}
		}
		if enc == nil {
			enc, mt = negotiate("")
		}
	}
	SetContentType(w, mt)
	return enc
}

// RequestEncoder returns a HTTP request encoder.
// The encoder uses package encoding/json.
func RequestEncoder(r *http.Request) Encoder {
	var buf bytes.Buffer
	r.Body = ioutil.NopCloser(&buf)
	return json.NewEncoder(&buf)
}

// ResponseDecoder returns a HTTP response decoder.
// The decoder handles the following content types:
//
//   * application/json using package encoding/json (default)
//   * application/xml using package encoding/xml
//   * application/gob using package encoding/gob
//   * text/html and text/plain for strings
//
func ResponseDecoder(resp *http.Response) Decoder {
	ct := resp.Header.Get("Content-Type")
	if ct == "" {
		return json.NewDecoder(resp.Body)
	}
	if mediaType, _, err := mime.ParseMediaType(ct); err == nil {
		ct = mediaType
	}
	switch {
	case ct == "application/json" || strings.HasSuffix(ct, "+json"):
		return json.NewDecoder(resp.Body)
	case ct == "application/xml" || strings.HasSuffix(ct, "+xml"):
		return xml.NewDecoder(resp.Body)
	case ct == "application/gob" || strings.HasSuffix(ct, "+gob"):
		return gob.NewDecoder(resp.Body)
	case ct == "text/html" || ct == "text/plain" ||
		strings.HasSuffix(ct, "+html") || strings.HasSuffix(ct, "+txt"):
		return newTextDecoder(resp.Body, ct)
	default:
		return json.NewDecoder(resp.Body)
	}
}

// ErrorEncoder returns an encoder that encodes errors returned by service
// methods. The encoder checks whether the error is a goa ServiceError struct
// and if so uses the error temporary and timeout fields to infer a proper HTTP
// status code and marshals the error struct to the body using the provided
// encoder. If the error is not a goa ServiceError struct then it is encoded
// as a permanent internal server error.
func ErrorEncoder(encoder func(context.Context, http.ResponseWriter) Encoder) func(context.Context, http.ResponseWriter, error) error {
	return func(ctx context.Context, w http.ResponseWriter, err error) error {
		enc := encoder(ctx, w)
		resp := NewErrorResponse(err)
		w.WriteHeader(resp.StatusCode())
		return enc.Encode(resp)
	}
}

// Decode implements the Decoder interface. It simply calls f(v).
func (f EncodingFunc) Decode(v interface{}) error { return f(v) }

// Encode implements the Encoder interface. It simply calls f(v).
func (f EncodingFunc) Encode(v interface{}) error { return f(v) }

// SetContentType initializes the response Content-Type header given a MIME
// type. If the Content-Type header is already set and the MIME type is
// "application/json" or "application/xml" then SetContentType appends a suffix
// to the header ("+json" or "+xml" respectively).
func SetContentType(w http.ResponseWriter, ct string) {
	h := w.Header().Get("Content-Type")
	if h == "" {
		w.Header().Set("Content-Type", ct)
		return
	}
	// RFC6839 only defines suffixes for a few mime types, we only concern
	// ourselves with JSON and XML.
	if ct != "application/json" && ct != "application/xml" {
		w.Header().Set("Content-Type", ct)
		return
	}
	if strings.Contains(h, "+") {
		return
	}
	suffix := "+json"
	if ct == "application/xml" {
		suffix = "+xml"
	}
	w.Header().Set("Content-Type", h+suffix)
}

func newTextEncoder(w io.Writer, ct string) Encoder {
	return &textEncoder{w, ct}
}

type textEncoder struct {
	w  io.Writer
	ct string
}

func (e *textEncoder) Encode(v interface{}) error {
	var err error

	switch c := v.(type) {
	case string:
		_, err = e.w.Write([]byte(c))
	case *string: // v may be a string pointer when the Response Body is set to the field of a custom response type.
		_, err = e.w.Write([]byte(*c))
	case []byte:
		_, err = e.w.Write(c)
	default:
		err = fmt.Errorf("can't encode %T as %s", c, e.ct)
	}

	return err
}

func newTextDecoder(r io.Reader, ct string) Decoder {
	return &textDecoder{r, ct}
}

type textDecoder struct {
	r  io.Reader
	ct string
}

func (e *textDecoder) Decode(v interface{}) error {
	b, err := ioutil.ReadAll(e.r)
	if err != nil {
		return err
	}

	switch c := v.(type) {
	case *string:
		*c = string(b)
	case *[]byte:
		*c = b
	default:
		err = fmt.Errorf("can't decode %s to %T", e.ct, c)
	}

	return err
}
<commit_msg>Support text/html and text/plain by RequestDecoder (#2310) (#2311)<commit_after>package http

import (
	"bytes"
	"context"
	"encoding/gob"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"mime"
	"net/http"
	"strings"
)

const (
	// AcceptTypeKey is the context key used to store the value of the HTTP
	// request Accept-Type header. The value may be used by encoders and
	// decoders to implement a content type negotiation algorithm.
	AcceptTypeKey contextKey = iota + 1
	// ContentTypeKey is the context key used to store the value of the HTTP
	// response Content-Type header when explicitly set in the DSL. The value
	// may be used by encoders to set the header appropriately.
	ContentTypeKey
)

type (
	// Decoder provides the actual decoding algorithm used to load HTTP
	// request and response bodies.
	Decoder interface {
		// Decode decodes into v.
		Decode(v interface{}) error
	}

	// Encoder provides the actual encoding algorithm used to write HTTP
	// request and response bodies.
	Encoder interface {
		// Encode encodes v.
		Encode(v interface{}) error
	}

	// EncodingFunc allows a function with appropriate signature to act as a
	// Decoder/Encoder.
	EncodingFunc func(v interface{}) error

	// private type used to define context keys.
	contextKey int
)

// RequestDecoder returns a HTTP request body decoder suitable for the given
// request. The decoder handles the following mime types:
//
//     * application/json using package encoding/json
//     * application/xml using package encoding/xml
//     * application/gob using package encoding/gob
//
// RequestDecoder defaults to the JSON decoder if the request "Content-Type"
// header does not match any of the supported mime type or is missing
// altogether.
func RequestDecoder(r *http.Request) Decoder {
	contentType := r.Header.Get("Content-Type")
	if contentType == "" {
		// default to JSON
		contentType = "application/json"
	} else {
		// sanitize
		if mediaType, _, err := mime.ParseMediaType(contentType); err == nil {
			contentType = mediaType
		}
	}
	switch contentType {
	case "application/json":
		return json.NewDecoder(r.Body)
	case "application/gob":
		return gob.NewDecoder(r.Body)
	case "application/xml":
		return xml.NewDecoder(r.Body)
	case "text/html", "text/plain":
		return newTextDecoder(r.Body, contentType)
	default:
		return json.NewDecoder(r.Body)
	}
}

// ResponseEncoder returns a HTTP response encoder leveraging the mime type
// set in the context under the AcceptTypeKey or the ContentTypeKey if any.
// The encoder supports the following mime types:
//
//     * application/json using package encoding/json
//     * application/xml using package encoding/xml
//     * application/gob using package encoding/gob
//     * text/html and text/plain for strings
//
// ResponseEncoder defaults to the JSON encoder if the context AcceptTypeKey or
// ContentTypeKey value does not match any of the supported mime types or is
// missing altogether.
func ResponseEncoder(ctx context.Context, w http.ResponseWriter) Encoder {
	negotiate := func(a string) (Encoder, string) {
		switch a {
		case "", "application/json":
			// default to JSON
			return json.NewEncoder(w), "application/json"
		case "application/xml":
			return xml.NewEncoder(w), "application/xml"
		case "application/gob":
			return gob.NewEncoder(w), "application/gob"
		case "text/html", "text/plain":
			return newTextEncoder(w, a), a
		}
		return nil, ""
	}
	var accept string
	{
		if a := ctx.Value(AcceptTypeKey); a != nil {
			accept = a.(string)
		}
	}
	var ct string
	{
		if a := ctx.Value(ContentTypeKey); a != nil {
			ct = a.(string)
		}
	}
	var (
		enc Encoder
		mt  string
		err error
	)
	{
		if ct != "" {
			// If content type explicitly set in the DSL, infer the response encoder
			// from the content type context key.
			if mt, _, err = mime.ParseMediaType(ct); err == nil {
				switch {
				case ct == "application/json" || strings.HasSuffix(ct, "+json"):
					enc = json.NewEncoder(w)
				case ct == "application/xml" || strings.HasSuffix(ct, "+xml"):
					enc = xml.NewEncoder(w)
				case ct == "application/gob" || strings.HasSuffix(ct, "+gob"):
					enc = gob.NewEncoder(w)
				case ct == "text/html" || ct == "text/plain" ||
					strings.HasSuffix(ct, "+html") || strings.HasSuffix(ct, "+txt"):
					enc = newTextEncoder(w, ct)
				default:
					enc = json.NewEncoder(w)
				}
			}
			SetContentType(w, mt)
			return enc
		}
		// If Accept header exists in the request, infer the response encoder
		// from the header value.
		if enc, mt = negotiate(accept); enc == nil {
			// attempt to normalize
			if mt, _, err = mime.ParseMediaType(accept); err == nil {
				enc, mt = negotiate(mt)
			}
		}
		if enc == nil {
			enc, mt = negotiate("")
		}
	}
	SetContentType(w, mt)
	return enc
}

// RequestEncoder returns a HTTP request encoder.
// The encoder uses package encoding/json.
func RequestEncoder(r *http.Request) Encoder {
	var buf bytes.Buffer
	r.Body = ioutil.NopCloser(&buf)
	return json.NewEncoder(&buf)
}

// ResponseDecoder returns a HTTP response decoder.
// The decoder handles the following content types:
//
//   * application/json using package encoding/json (default)
//   * application/xml using package encoding/xml
//   * application/gob using package encoding/gob
//   * text/html and text/plain for strings
//
func ResponseDecoder(resp *http.Response) Decoder {
	ct := resp.Header.Get("Content-Type")
	if ct == "" {
		return json.NewDecoder(resp.Body)
	}
	if mediaType, _, err := mime.ParseMediaType(ct); err == nil {
		ct = mediaType
	}
	switch {
	case ct == "application/json" || strings.HasSuffix(ct, "+json"):
		return json.NewDecoder(resp.Body)
	case ct == "application/xml" || strings.HasSuffix(ct, "+xml"):
		return xml.NewDecoder(resp.Body)
	case ct == "application/gob" || strings.HasSuffix(ct, "+gob"):
		return gob.NewDecoder(resp.Body)
	case ct == "text/html" || ct == "text/plain" ||
		strings.HasSuffix(ct, "+html") || strings.HasSuffix(ct, "+txt"):
		return newTextDecoder(resp.Body, ct)
	default:
		return json.NewDecoder(resp.Body)
	}
}

// ErrorEncoder returns an encoder that encodes errors returned by service
// methods. The encoder checks whether the error is a goa ServiceError struct
// and if so uses the error temporary and timeout fields to infer a proper HTTP
// status code and marshals the error struct to the body using the provided
// encoder. If the error is not a goa ServiceError struct then it is encoded
// as a permanent internal server error.
func ErrorEncoder(encoder func(context.Context, http.ResponseWriter) Encoder) func(context.Context, http.ResponseWriter, error) error {
	return func(ctx context.Context, w http.ResponseWriter, err error) error {
		enc := encoder(ctx, w)
		resp := NewErrorResponse(err)
		w.WriteHeader(resp.StatusCode())
		return enc.Encode(resp)
	}
}

// Decode implements the Decoder interface. It simply calls f(v).
func (f EncodingFunc) Decode(v interface{}) error { return f(v) }

// Encode implements the Encoder interface. It simply calls f(v).
func (f EncodingFunc) Encode(v interface{}) error { return f(v) }

// SetContentType initializes the response Content-Type header given a MIME
// type. If the Content-Type header is already set and the MIME type is
// "application/json" or "application/xml" then SetContentType appends a suffix
// to the header ("+json" or "+xml" respectively).
func SetContentType(w http.ResponseWriter, ct string) {
	h := w.Header().Get("Content-Type")
	if h == "" {
		w.Header().Set("Content-Type", ct)
		return
	}
	// RFC6839 only defines suffixes for a few mime types, we only concern
	// ourselves with JSON and XML.
	if ct != "application/json" && ct != "application/xml" {
		w.Header().Set("Content-Type", ct)
		return
	}
	if strings.Contains(h, "+") {
		return
	}
	suffix := "+json"
	if ct == "application/xml" {
		suffix = "+xml"
	}
	w.Header().Set("Content-Type", h+suffix)
}

func newTextEncoder(w io.Writer, ct string) Encoder {
	return &textEncoder{w, ct}
}

type textEncoder struct {
	w  io.Writer
	ct string
}

func (e *textEncoder) Encode(v interface{}) error {
	var err error

	switch c := v.(type) {
	case string:
		_, err = e.w.Write([]byte(c))
	case *string: // v may be a string pointer when the Response Body is set to the field of a custom response type.
		_, err = e.w.Write([]byte(*c))
	case []byte:
		_, err = e.w.Write(c)
	default:
		err = fmt.Errorf("can't encode %T as %s", c, e.ct)
	}

	return err
}

func newTextDecoder(r io.Reader, ct string) Decoder {
	return &textDecoder{r, ct}
}

type textDecoder struct {
	r  io.Reader
	ct string
}

func (e *textDecoder) Decode(v interface{}) error {
	b, err := ioutil.ReadAll(e.r)
	if err != nil {
		return err
	}

	switch c := v.(type) {
	case *string:
		*c = string(b)
	case *[]byte:
		*c = b
	default:
		err = fmt.Errorf("can't decode %s to %T", e.ct, c)
	}

	return err
}
<|endoftext|>


<commit_before>package snmpclient2

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strings"
	"sync"

	"github.com/runner-mei/snmpclient2/asn1"
)

type OidAndValue struct {
	Oid   Oid
	Value Variable
}

func compareOidAdValue(a1, b1 Item) int {
	a, ok := a1.(Oid)
	if !ok {
		v, ok := a1.(*OidAndValue)
		if !ok {
			panic("a1 is not OidAndValue")
		}
		a = v.Oid
	}

	b, ok := b1.(Oid)
	if !ok {
		v, ok := b1.(*OidAndValue)
		if !ok {
			panic("b1 is not OidAndValue")
		}
		b = v.Oid
	}

	a_uint32s := a.Value
	b_uint32s := b.Value
	for idx, c := range a_uint32s {
		if idx >= len(b_uint32s) {
			return 1
		}
		if c == b_uint32s[idx] {
			continue
		}

		if c < b_uint32s[idx] {
			return -1
		}
		return 1
	}
	if len(a_uint32s) == len(b_uint32s) {
		return 0
	}
	return -1
}

func convertOid(a interface{}) *Oid {
	o, ok := a.(Oid)
	if ok {
		return &o
	}
	v, ok := a.(*OidAndValue)
	if ok {
		return &v.Oid
	}
	panic("a1 is not OidAndValue")
}

func compareOidAdValueWith(a1, b1 Item) int {
	r := compareOidAdValue(a1, b1)
	if r > 0 {
		fmt.Println(convertOid(a1).String() + ">" + convertOid(b1).String())
	} else if r < 0 {
		fmt.Println(convertOid(a1).String() + "<" + convertOid(b1).String())
	} else {
		fmt.Println(convertOid(a1).String() + "==" + convertOid(b1).String())
	}

	return r
}

func NewMibTree() *Tree {
	return NewTree(compareOidAdValue)
}

// ******************************************
//  It is for test.
type UdpServer struct {
	name       string
	origin     string
	conn       net.PacketConn
	listenAddr net.Addr
	waitGroup  sync.WaitGroup
	mpv1       Security
	//priv_type  PrivType
	//priv_key []byte

	return_error_if_oid_not_exists bool
	is_update_mibs                 bool
	mibs                           *Tree
}

func NewUdpServerFromFile(nm, addr, file string, is_update_mibs bool) (*UdpServer, error) {
	srv := &UdpServer{name: nm,
		origin:         addr,
		is_update_mibs: is_update_mibs,
		mibs:           NewMibTree(),
		mpv1:           NewCommunity()}
	r, e := os.Open(file)
	if nil != e {
		return nil, e
	}
	if e := Read(r, func(oid Oid, value Variable) error {
		if ok := srv.mibs.Insert(&OidAndValue{Oid: oid,
			Value: value}); !ok {
			if srv.is_update_mibs {
				if ok = srv.mibs.DeleteWithKey(oid); !ok {
					return errors.New("insert '" + oid.String() + "' failed, delete failed.")
				}
				if ok = srv.mibs.Insert(&OidAndValue{Oid: oid,
					Value: value}); !ok {
					return errors.New("insert '" + oid.String() + "' failed.")
				}
			} else {
				return errors.New("insert '" + oid.String() + "' failed.")
			}
		}
		return nil
	}); nil != e {
		return nil, e
	}
	return srv, srv.start()
}

func NewUdpServerFromString(nm, addr, mibs string, is_update_mibs bool) (*UdpServer, error) {
	srv := &UdpServer{name: nm,
		origin:         addr,
		is_update_mibs: is_update_mibs,
		mibs:           NewMibTree(),
		mpv1:           NewCommunity()}
	if e := Read(bytes.NewReader([]byte(mibs)), func(oid Oid, value Variable) error {
		if ok := srv.mibs.Insert(&OidAndValue{Oid: oid,
			Value: value}); !ok {

			if srv.is_update_mibs {
				if ok = srv.mibs.DeleteWithKey(oid); !ok {
					return errors.New("insert '" + oid.String() + "' failed, delete failed.")
				}
				if ok = srv.mibs.Insert(&OidAndValue{Oid: oid,
					Value: value}); !ok {
					return errors.New("insert '" + oid.String() + "' failed.")
				}
			} else {
				return errors.New("insert '" + oid.String() + "' failed.")
			}
		}
		return nil
	}); nil != e {
		return nil, e
	}
	return srv, srv.start()
}

func (self *UdpServer) ReturnErrorIfOidNotExists(status bool) *UdpServer {
	self.return_error_if_oid_not_exists = status
	return self
}

func (self *UdpServer) ReloadMibsFromFile(file string) error {
	bs, e := ioutil.ReadFile(file)
	if nil != e {
		return e
	}

	self.mibs = NewMibTree()
	return self.LoadMibsFromString(string(bs))
}

func (self *UdpServer) ReloadMibsFromString(mibs string) error {
	self.mibs = NewMibTree()
	return self.LoadMibsFromString(mibs)
}

func (self *UdpServer) LoadFile(file string) error {
	mibs, e := ioutil.ReadFile(file)
	if nil != e {
		return e
	}
	return self.LoadMibsFromString(string(mibs))
}

func (self *UdpServer) LoadMibsFromString(mibs string) error {
	if e := Read(bytes.NewReader([]byte(mibs)), func(oid Oid, value Variable) error {
		if ok := self.mibs.Insert(&OidAndValue{Oid: oid,
			Value: value}); !ok {

			if self.is_update_mibs {
				if ok = self.mibs.DeleteWithKey(oid); !ok {
					return errors.New("insert '" + oid.String() + "' failed, delete failed.")
				}
				if ok = self.mibs.Insert(&OidAndValue{Oid: oid,
					Value: value}); !ok {
					return errors.New("insert '" + oid.String() + "' failed.")
				}
			} else {
				return errors.New("insert '" + oid.String() + "' failed.")
			}
		}
		return nil
	}); nil != e {
		return e
	}
	return nil
}
func (self *UdpServer) GetPort() string {
	s := self.listenAddr.String()
	if i := strings.LastIndex(s, ":"); -1 != i {
		return s[i+1:]
	}
	return ""
}

func (self *UdpServer) Close() {
	self.conn.Close()
	self.waitGroup.Wait()
}

func (self *UdpServer) start() error {
	var conn net.PacketConn
	var e error

	if nil == self.listenAddr {
		conn, e = net.ListenPacket("udp", self.origin)
	} else {
		conn, e = net.ListenPacket("udp", self.listenAddr.String())
	}
	if nil != e {
		return e
	}

	self.conn = conn
	self.listenAddr = conn.LocalAddr()

	self.waitGroup.Add(1)
	go self.serve()

	return nil
}

func (self *UdpServer) serve() {
	defer func() {
		self.conn = nil
		self.waitGroup.Done()
	}()

	var cached_bytes [10240]byte

	for {
		n, addr, err := self.conn.ReadFrom(cached_bytes[:])
		if nil != err {
			log.Println("[", self.name, "]", err.Error())
			break
		}

		func(recv_bytes []byte) {
			var raw asn1.RawValue
			_, err = asn1.Unmarshal(recv_bytes, &raw)
			if err != nil {
				log.Printf("["+self.name+"]Invalid MessageV3 object - %s : [%s]",
					err.Error(), ToHexStr(recv_bytes, " "))
				return
			}

			if raw.Class != asn1.ClassUniversal || raw.Tag != asn1.TagSequence || !raw.IsCompound {
				log.Printf("["+self.name+"]Invalid MessageV3 object - Class [%02x], Tag [%02x] : [%s]",
					raw.FullBytes[0], raw.Tag, ToHexStr(recv_bytes, " "))
				return
			}
			next := raw.Bytes

			var version int
			next, err = asn1.Unmarshal(next, &version)
			if err != nil {
				log.Printf("["+self.name+"]Invalid MessageV3 object - %s : [%s]",
					err.Error(), ToHexStr(recv_bytes, " "))
				return
			}

			if SnmpVersion(version) == V3 {
				// var managedId int
				// next, err = asn1.Unmarshal(next, &managedId)
				// if err != nil {
				// 	return
				// }

				// pdu := &ScopedPdu{}
				// recvMsg := NewMessage(SnmpVersion(version), pdu)
				// _, err = recvMsg.Unmarshal(recv_bytes)
				// if err != nil {
				// 	log.Printf(client.logCtx, "Failed to Unmarshal message - %s : [%s]",
				// 		err.Error(), ToHexStr(recv_bytes, " "))
				// 	return
				// }

				// var ok bool
				// request, ok = client.pendings[managedId]
				// if !ok {
				// 	log.Printf(client.logCtx, "request with requestId was", managedId, "is not exists.")
				// 	return
				// }
				// delete(client.pendings, managedId)
				// request.response = pdu

				// err = client.mpv3.ProcessIncomingMessage(snmp, nil, recvMsg)
				// if err != nil {
				// 	log.Printf(client.logCtx, "Failed to process incoming message - %s : [%s]",
				// 		err.Error(), ToHexStr(recv_bytes, " "))
				// 	return
				// }

				log.Printf("["+self.name+"]Failed to process incoming message - v3 message is unsupported : [%s]",
					ToHexStr(recv_bytes, " "))
				return
			}
			recvMsg := &MessageV1{
				version: SnmpVersion(version),
				pdu:     &PduV1{},
			}
			_, err = recvMsg.Unmarshal(recv_bytes)
			if err != nil {
				log.Printf("["+self.name+"]Failed to Unmarshal message - %s : [%s]",
					err.Error(), ToHexStr(recv_bytes, " "))
				return
			}

			err = self.mpv1.ProcessIncomingMessage(nil, recvMsg)
			if err != nil {
				log.Printf("["+self.name+"]Failed to process incoming message - %s : [%s]",
					err.Error(), ToHexStr(recv_bytes, " "))
				return
			}
			self.on_v2(addr, recvMsg, recv_bytes)

		}(cached_bytes[:n])
	}
}

func (self *UdpServer) on_v2(addr net.Addr, p *MessageV1, cached_bytes []byte) {
	pdu := &PduV1{
		pduType:   GetResponse,
		requestId: p.PDU().RequestId(),
	}

	res := &MessageV1{
		version: p.Version(),
		pdu:     pdu,
	}
	//res.SetMaxMsgSize(p.GetMaxMsgSize())

	switch p.PDU().PduType() {
	case GetRequest:
		for _, vb := range p.PDU().VariableBindings() {
			v := self.GetValueByOid(vb.Oid)
			if nil == v {
				if self.return_error_if_oid_not_exists {
					pdu.SetErrorStatus(NoSuchName)
					break
				}
				continue
			}
			res.PDU().AppendVariableBinding(vb.Oid, v)
		}
	case GetNextRequest:
		for _, vb := range p.PDU().VariableBindings() {
			o, v := self.GetNextValueByOid(vb.Oid)
			if nil == v {
				continue
			}
			res.PDU().AppendVariableBinding(*o, v)
		}
	default:
		log.Println("[", self.name, "] snmp type is not supported.")
	}

	err := NewCommunity().GenerateRequestMessage(&Arguments{Community: ""}, res)
	if err != nil {
		log.Println("[", self.name, "] failed to generate request,", err)
		return
	}

	s, err := res.Marshal()
	if err != nil {
		log.Println("[", self.name, "] failed to marshal,", err)
		return
	}
	if _, e := self.conn.WriteTo(s, addr); nil != e {
		log.Println("[", self.name, "] failed to write response,", e)
		return
	}
}

func (self *UdpServer) GetValueByOid(oid Oid) Variable {
	if v := self.mibs.Get(oid); nil != v {
		if sv, ok := v.(*OidAndValue); ok {
			return sv.Value
		}
		panic(fmt.Sprintf("it is not a Variable - [%T]%v", v, v))
	}
	return nil
}

func (self *UdpServer) GetNextValueByOid(oid Oid) (*Oid, Variable) {
	it := self.mibs.FindGE(oid)
	if it.Limit() {
		return nil, nil
	}
	v := it.Item()
	if nil == v {
		return nil, nil
	}
	sv, ok := v.(*OidAndValue)
	if !ok {
		panic(fmt.Sprintf("it is not a Variable - [%T]%v", v, v))
	}

	if 0 != compareOidAdValue(oid, sv.Oid) {
		return &sv.Oid, sv.Value
	}
	it = it.Next()
	if it.Limit() {
		return nil, nil
	}
	v = it.Item()
	if nil == v {
		return nil, nil
	}
	sv, ok = v.(*OidAndValue)
	if ok {
		return &sv.Oid, sv.Value
	}
	panic(fmt.Sprintf("it is not a Variable - [%T]%v", v, v))
}
<commit_msg>add UdpServer.Restart()<commit_after>package snmpclient2

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strings"
	"sync"

	"github.com/runner-mei/snmpclient2/asn1"
)

type OidAndValue struct {
	Oid   Oid
	Value Variable
}

func compareOidAdValue(a1, b1 Item) int {
	a, ok := a1.(Oid)
	if !ok {
		v, ok := a1.(*OidAndValue)
		if !ok {
			panic("a1 is not OidAndValue")
		}
		a = v.Oid
	}

	b, ok := b1.(Oid)
	if !ok {
		v, ok := b1.(*OidAndValue)
		if !ok {
			panic("b1 is not OidAndValue")
		}
		b = v.Oid
	}

	a_uint32s := a.Value
	b_uint32s := b.Value
	for idx, c := range a_uint32s {
		if idx >= len(b_uint32s) {
			return 1
		}
		if c == b_uint32s[idx] {
			continue
		}

		if c < b_uint32s[idx] {
			return -1
		}
		return 1
	}
	if len(a_uint32s) == len(b_uint32s) {
		return 0
	}
	return -1
}

func convertOid(a interface{}) *Oid {
	o, ok := a.(Oid)
	if ok {
		return &o
	}
	v, ok := a.(*OidAndValue)
	if ok {
		return &v.Oid
	}
	panic("a1 is not OidAndValue")
}

func compareOidAdValueWith(a1, b1 Item) int {
	r := compareOidAdValue(a1, b1)
	if r > 0 {
		fmt.Println(convertOid(a1).String() + ">" + convertOid(b1).String())
	} else if r < 0 {
		fmt.Println(convertOid(a1).String() + "<" + convertOid(b1).String())
	} else {
		fmt.Println(convertOid(a1).String() + "==" + convertOid(b1).String())
	}

	return r
}

func NewMibTree() *Tree {
	return NewTree(compareOidAdValue)
}

// ******************************************
//  It is for test.
type UdpServer struct {
	name       string
	origin     string
	conn       net.PacketConn
	listenAddr net.Addr
	waitGroup  sync.WaitGroup
	mpv1       Security
	//priv_type  PrivType
	//priv_key []byte

	return_error_if_oid_not_exists bool
	is_update_mibs                 bool
	mibs                           *Tree
}

func NewUdpServerFromFile(nm, addr, file string, is_update_mibs bool) (*UdpServer, error) {
	srv := &UdpServer{name: nm,
		origin:         addr,
		is_update_mibs: is_update_mibs,
		mibs:           NewMibTree(),
		mpv1:           NewCommunity()}
	r, e := os.Open(file)
	if nil != e {
		return nil, e
	}
	if e := Read(r, func(oid Oid, value Variable) error {
		if ok := srv.mibs.Insert(&OidAndValue{Oid: oid,
			Value: value}); !ok {
			if srv.is_update_mibs {
				if ok = srv.mibs.DeleteWithKey(oid); !ok {
					return errors.New("insert '" + oid.String() + "' failed, delete failed.")
				}
				if ok = srv.mibs.Insert(&OidAndValue{Oid: oid,
					Value: value}); !ok {
					return errors.New("insert '" + oid.String() + "' failed.")
				}
			} else {
				return errors.New("insert '" + oid.String() + "' failed.")
			}
		}
		return nil
	}); nil != e {
		return nil, e
	}
	return srv, srv.start()
}

func NewUdpServerFromString(nm, addr, mibs string, is_update_mibs bool) (*UdpServer, error) {
	srv := &UdpServer{name: nm,
		origin:         addr,
		is_update_mibs: is_update_mibs,
		mibs:           NewMibTree(),
		mpv1:           NewCommunity()}
	if e := Read(bytes.NewReader([]byte(mibs)), func(oid Oid, value Variable) error {
		if ok := srv.mibs.Insert(&OidAndValue{Oid: oid,
			Value: value}); !ok {

			if srv.is_update_mibs {
				if ok = srv.mibs.DeleteWithKey(oid); !ok {
					return errors.New("insert '" + oid.String() + "' failed, delete failed.")
				}
				if ok = srv.mibs.Insert(&OidAndValue{Oid: oid,
					Value: value}); !ok {
					return errors.New("insert '" + oid.String() + "' failed.")
				}
			} else {
				return errors.New("insert '" + oid.String() + "' failed.")
			}
		}
		return nil
	}); nil != e {
		return nil, e
	}
	return srv, srv.start()
}

func (self *UdpServer) ReturnErrorIfOidNotExists(status bool) *UdpServer {
	self.return_error_if_oid_not_exists = status
	return self
}

func (self *UdpServer) ReloadMibsFromFile(file string) error {
	bs, e := ioutil.ReadFile(file)
	if nil != e {
		return e
	}

	self.mibs = NewMibTree()
	return self.LoadMibsFromString(string(bs))
}

func (self *UdpServer) ReloadMibsFromString(mibs string) error {
	self.mibs = NewMibTree()
	return self.LoadMibsFromString(mibs)
}

func (self *UdpServer) LoadFile(file string) error {
	mibs, e := ioutil.ReadFile(file)
	if nil != e {
		return e
	}
	return self.LoadMibsFromString(string(mibs))
}

func (self *UdpServer) LoadMibsFromString(mibs string) error {
	if e := Read(bytes.NewReader([]byte(mibs)), func(oid Oid, value Variable) error {
		if ok := self.mibs.Insert(&OidAndValue{Oid: oid,
			Value: value}); !ok {

			if self.is_update_mibs {
				if ok = self.mibs.DeleteWithKey(oid); !ok {
					return errors.New("insert '" + oid.String() + "' failed, delete failed.")
				}
				if ok = self.mibs.Insert(&OidAndValue{Oid: oid,
					Value: value}); !ok {
					return errors.New("insert '" + oid.String() + "' failed.")
				}
			} else {
				return errors.New("insert '" + oid.String() + "' failed.")
			}
		}
		return nil
	}); nil != e {
		return e
	}
	return nil
}
func (self *UdpServer) GetPort() string {
	s := self.listenAddr.String()
	if i := strings.LastIndex(s, ":"); -1 != i {
		return s[i+1:]
	}
	return ""
}

func (self *UdpServer) Close() {
	self.conn.Close()
	self.waitGroup.Wait()
	self.conn = nil
}

func (self *UdpServer) Restart() error {
	self.Close()

	self.listenAddr = nil
	log.Println("udp server is exited")
	err := self.start()
	if err == nil {

		log.Println("udp server is restarted, listen at", self.listenAddr.String())
	}
	return err
}

func (self *UdpServer) start() error {
	var conn net.PacketConn
	var e error

	if nil == self.listenAddr {
		conn, e = net.ListenPacket("udp", self.origin)
	} else {
		conn, e = net.ListenPacket("udp", self.listenAddr.String())
	}
	if nil != e {
		return e
	}

	self.conn = conn
	self.listenAddr = conn.LocalAddr()

	self.waitGroup.Add(1)
	go self.serve()

	return nil
}

func (self *UdpServer) serve() {
	defer func() {
		self.conn = nil
		self.waitGroup.Done()
	}()

	var cached_bytes [10240]byte

	for {
		n, addr, err := self.conn.ReadFrom(cached_bytes[:])
		if nil != err {
			log.Println("[", self.name, "]", err.Error())
			break
		}

		func(recv_bytes []byte) {
			var raw asn1.RawValue
			_, err = asn1.Unmarshal(recv_bytes, &raw)
			if err != nil {
				log.Printf("["+self.name+"]Invalid MessageV3 object - %s : [%s]",
					err.Error(), ToHexStr(recv_bytes, " "))
				return
			}

			if raw.Class != asn1.ClassUniversal || raw.Tag != asn1.TagSequence || !raw.IsCompound {
				log.Printf("["+self.name+"]Invalid MessageV3 object - Class [%02x], Tag [%02x] : [%s]",
					raw.FullBytes[0], raw.Tag, ToHexStr(recv_bytes, " "))
				return
			}
			next := raw.Bytes

			var version int
			next, err = asn1.Unmarshal(next, &version)
			if err != nil {
				log.Printf("["+self.name+"]Invalid MessageV3 object - %s : [%s]",
					err.Error(), ToHexStr(recv_bytes, " "))
				return
			}

			if SnmpVersion(version) == V3 {
				// var managedId int
				// next, err = asn1.Unmarshal(next, &managedId)
				// if err != nil {
				// 	return
				// }

				// pdu := &ScopedPdu{}
				// recvMsg := NewMessage(SnmpVersion(version), pdu)
				// _, err = recvMsg.Unmarshal(recv_bytes)
				// if err != nil {
				// 	log.Printf(client.logCtx, "Failed to Unmarshal message - %s : [%s]",
				// 		err.Error(), ToHexStr(recv_bytes, " "))
				// 	return
				// }

				// var ok bool
				// request, ok = client.pendings[managedId]
				// if !ok {
				// 	log.Printf(client.logCtx, "request with requestId was", managedId, "is not exists.")
				// 	return
				// }
				// delete(client.pendings, managedId)
				// request.response = pdu

				// err = client.mpv3.ProcessIncomingMessage(snmp, nil, recvMsg)
				// if err != nil {
				// 	log.Printf(client.logCtx, "Failed to process incoming message - %s : [%s]",
				// 		err.Error(), ToHexStr(recv_bytes, " "))
				// 	return
				// }

				log.Printf("["+self.name+"]Failed to process incoming message - v3 message is unsupported : [%s]",
					ToHexStr(recv_bytes, " "))
				return
			}
			recvMsg := &MessageV1{
				version: SnmpVersion(version),
				pdu:     &PduV1{},
			}
			_, err = recvMsg.Unmarshal(recv_bytes)
			if err != nil {
				log.Printf("["+self.name+"]Failed to Unmarshal message - %s : [%s]",
					err.Error(), ToHexStr(recv_bytes, " "))
				return
			}

			err = self.mpv1.ProcessIncomingMessage(nil, recvMsg)
			if err != nil {
				log.Printf("["+self.name+"]Failed to process incoming message - %s : [%s]",
					err.Error(), ToHexStr(recv_bytes, " "))
				return
			}
			self.on_v2(addr, recvMsg, recv_bytes)

		}(cached_bytes[:n])
	}
}

func (self *UdpServer) on_v2(addr net.Addr, p *MessageV1, cached_bytes []byte) {
	pdu := &PduV1{
		pduType:   GetResponse,
		requestId: p.PDU().RequestId(),
	}

	res := &MessageV1{
		version: p.Version(),
		pdu:     pdu,
	}
	//res.SetMaxMsgSize(p.GetMaxMsgSize())

	switch p.PDU().PduType() {
	case GetRequest:
		for _, vb := range p.PDU().VariableBindings() {
			v := self.GetValueByOid(vb.Oid)
			if nil == v {
				if self.return_error_if_oid_not_exists {
					pdu.SetErrorStatus(NoSuchName)
					break
				}
				continue
			}
			res.PDU().AppendVariableBinding(vb.Oid, v)
		}
	case GetNextRequest:
		for _, vb := range p.PDU().VariableBindings() {
			o, v := self.GetNextValueByOid(vb.Oid)
			if nil == v {
				continue
			}
			res.PDU().AppendVariableBinding(*o, v)
		}
	default:
		log.Println("[", self.name, "] snmp type is not supported.")
	}

	err := NewCommunity().GenerateRequestMessage(&Arguments{Community: ""}, res)
	if err != nil {
		log.Println("[", self.name, "] failed to generate request,", err)
		return
	}

	s, err := res.Marshal()
	if err != nil {
		log.Println("[", self.name, "] failed to marshal,", err)
		return
	}
	if _, e := self.conn.WriteTo(s, addr); nil != e {
		log.Println("[", self.name, "] failed to write response,", e)
		return
	}
}

func (self *UdpServer) GetValueByOid(oid Oid) Variable {
	if v := self.mibs.Get(oid); nil != v {
		if sv, ok := v.(*OidAndValue); ok {
			return sv.Value
		}
		panic(fmt.Sprintf("it is not a Variable - [%T]%v", v, v))
	}
	return nil
}

func (self *UdpServer) GetNextValueByOid(oid Oid) (*Oid, Variable) {
	it := self.mibs.FindGE(oid)
	if it.Limit() {
		return nil, nil
	}
	v := it.Item()
	if nil == v {
		return nil, nil
	}
	sv, ok := v.(*OidAndValue)
	if !ok {
		panic(fmt.Sprintf("it is not a Variable - [%T]%v", v, v))
	}

	if 0 != compareOidAdValue(oid, sv.Oid) {
		return &sv.Oid, sv.Value
	}
	it = it.Next()
	if it.Limit() {
		return nil, nil
	}
	v = it.Item()
	if nil == v {
		return nil, nil
	}
	sv, ok = v.(*OidAndValue)
	if ok {
		return &sv.Oid, sv.Value
	}
	panic(fmt.Sprintf("it is not a Variable - [%T]%v", v, v))
}
<|endoftext|>


<commit_before>/*
The MIT License (MIT)

Copyright (c) 2016 Ivan Dejanovic

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

package main

import (
	"fmt"
	"mlpl/cfg"
	"mlpl/lexer"
	"mlpl/parse"
	"mlpl/types"
	"os"
)

func main() {
	args := os.Args[1:]
	argc := len(args)

	if argc < 1 || argc > 2 {
		fmt.Println("Usage: <codefilename> [configurationfilename].")
		return
	}

	var reserved []types.ReservedWord

	if argc == 2 {
		reserved = cfg.GetConfigReservedWords(args[1])
	} else {
		reserved = cfg.GetDefaultReserved()
	}

	fmt.Println(reserved)

	tokens := parse.Parse(args[0], reserved)

	fmt.Println(tokens)

	treeNode := lexer.Lex(tokens)

	fmt.Println(treeNode)
}
<commit_msg>Fixed comment.<commit_after>/*
The MIT License (MIT)

Copyright (c) 2016 Ivan Dejanovic

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

package main

import (
	"fmt"
	"mlpl/cfg"
	"mlpl/lexer"
	"mlpl/parse"
	"mlpl/types"
	"os"
)

func main() {
	args := os.Args[1:]
	argc := len(args)

	if argc < 1 || argc > 2 {
		fmt.Println("Usage: <codefilename> [configurationfilename]")
		return
	}

	var reserved []types.ReservedWord

	if argc == 2 {
		reserved = cfg.GetConfigReservedWords(args[1])
	} else {
		reserved = cfg.GetDefaultReserved()
	}

	fmt.Println(reserved)

	tokens := parse.Parse(args[0], reserved)

	fmt.Println(tokens)

	treeNode := lexer.Lex(tokens)

	fmt.Println(treeNode)
}
<|endoftext|>


<commit_before>package gow

import (
	"io"
	"log"
	"net/http"
	"net"

	"github.com/gorilla/websocket"
)

// headers to drop
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",

	// added by me
	"Sec-Websocket-Accept",
}

type BackendSelector interface {
	Select(requestHost string) (string, error)
}

func ListenAndServeHTTP(address string, sel BackendSelector) error {
	proxyHandler := http.HandlerFunc(makeProxyHandlerFunc(sel))
	return http.ListenAndServe(address, proxyHandler)
}

func makeProxyHandlerFunc(sel BackendSelector) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		backend, err := sel.Select(r.Host)

		if err == nil {
			proxyRequest(w, r, backend)
		} else {
			crash, ok := err.(BootCrash)
			if ok {
				w.Header()["Content-Type"] = []string{"text/html"}
				w.WriteHeader(500)

				w.Write([]byte("
App crashed during boot :(

"))
				crash.Log.WriteTo(w)
				w.Write([]byte("
"))
				w.Write([]byte(terminalFormattingPostamble))
			} else {
				w.WriteHeader(500)
				w.Write([]byte("Failed to spawn backend: "))
				w.Write([]byte(err.Error()))
			}
		}
	}
}

func proxyRequest(w http.ResponseWriter, r *http.Request, backendAddress string) {
	r.RequestURI = ""

	if r.Header["Connection"] != nil && r.Header["Connection"][0] == "Upgrade" &&
		r.Header["Upgrade"] != nil && r.Header["Upgrade"][0] == "websocket" {
		log.Println("websocket connection requested by client; doing things")
		proxyWebsocket(w, r, backendAddress)
		return
	}

	r.URL.Scheme = "http"
	r.URL.Host = backendAddress

	resp, err := http.DefaultTransport.RoundTrip(r)

	if err != nil {
		log.Println(err)
		w.WriteHeader(502)
		w.Write([]byte{})
		return
	}

	writeResponseHeader(w, resp)

	// just stream the body to the client
	_, err = io.Copy(w, resp.Body)
	if err != nil {
		log.Println(err)
	}
}

func proxyWebsocket(w http.ResponseWriter, clientRequest *http.Request, backendAddress string) {
	clientRequest.URL.Scheme = "ws"

	log.Println("setting it")

	f := func(theNet, addr string) (net.Conn, error) {
		return net.Dial(theNet, backendAddress)
	}

	log.Println("made literal")
	d := websocket.Dialer{NetDial: f}

	log.Println("set it")

	upstream_conn, resp, err := d.Dial(clientRequest.URL.String(), clientRequest.Header)
	if err != nil {
		log.Println(err, resp)
		w.WriteHeader(502)
		w.Write([]byte{})
		return
	}

	for k := range hopHeaders {
		delete(resp.Header, hopHeaders[k])
	}

	client_conn, err := websocket.Upgrade(w, clientRequest, resp.Header, 4096, 4096)
	if err != nil {
		log.Println(err)
		w.WriteHeader(502)
		w.Write([]byte{})
		return
	}

	go func() {
		for {
			messageType, p, err := client_conn.ReadMessage()
			if err != nil {
				log.Println("error while reading from client:", err)
				break
			}
			if err = upstream_conn.WriteMessage(messageType, p); err != nil {
				log.Println("error while writing to upstream:", err)
				break
			}
		}
		upstream_conn.Close()
	}()

	go func() {
		for {
			messageType, p, err := upstream_conn.ReadMessage()
			if err != nil {
				log.Println("error while reading from upstream:", err)
				break
			}
			if err = client_conn.WriteMessage(messageType, p); err != nil {
				log.Println("error while writing to client:", err)
				break
			}
		}

		client_conn.Close()
	}()
}

func writeResponseHeader(w http.ResponseWriter, r *http.Response) {
	for k := range r.Header {
		should_drop := false
		for i := range hopHeaders {
			if k == hopHeaders[i] {
				should_drop = true
				break
			}
		}

		if !should_drop {
			w.Header()[k] = r.Header[k]
		}
	}

	w.Header().Set("X-Forwarded-For", "127.0.0.1")
	w.WriteHeader(r.StatusCode)
}
<commit_msg>remove spew<commit_after>package gow

import (
	"io"
	"log"
	"net/http"
	"net"

	"github.com/gorilla/websocket"
)

// headers to drop
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",

	// added by me
	"Sec-Websocket-Accept",
}

type BackendSelector interface {
	Select(requestHost string) (string, error)
}

func ListenAndServeHTTP(address string, sel BackendSelector) error {
	proxyHandler := http.HandlerFunc(makeProxyHandlerFunc(sel))
	return http.ListenAndServe(address, proxyHandler)
}

func makeProxyHandlerFunc(sel BackendSelector) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		backend, err := sel.Select(r.Host)

		if err == nil {
			proxyRequest(w, r, backend)
		} else {
			crash, ok := err.(BootCrash)
			if ok {
				w.Header()["Content-Type"] = []string{"text/html"}
				w.WriteHeader(500)

				w.Write([]byte("
App crashed during boot :(

"))
				crash.Log.WriteTo(w)
				w.Write([]byte("
"))
				w.Write([]byte(terminalFormattingPostamble))
			} else {
				w.WriteHeader(500)
				w.Write([]byte("Failed to spawn backend: "))
				w.Write([]byte(err.Error()))
			}
		}
	}
}

func proxyRequest(w http.ResponseWriter, r *http.Request, backendAddress string) {
	r.RequestURI = ""

	if r.Header["Connection"] != nil && r.Header["Connection"][0] == "Upgrade" &&
		r.Header["Upgrade"] != nil && r.Header["Upgrade"][0] == "websocket" {
		proxyWebsocket(w, r, backendAddress)
		return
	}

	r.URL.Scheme = "http"
	r.URL.Host = backendAddress

	resp, err := http.DefaultTransport.RoundTrip(r)

	if err != nil {
		log.Println(err)
		w.WriteHeader(502)
		w.Write([]byte{})
		return
	}

	writeResponseHeader(w, resp)

	// just stream the body to the client
	_, err = io.Copy(w, resp.Body)
	if err != nil {
		log.Println(err)
	}
}

func proxyWebsocket(w http.ResponseWriter, clientRequest *http.Request, backendAddress string) {
	clientRequest.URL.Scheme = "ws"

	f := func(theNet, addr string) (net.Conn, error) {
		return net.Dial(theNet, backendAddress)
	}

	d := websocket.Dialer{NetDial: f}

	upstream_conn, resp, err := d.Dial(clientRequest.URL.String(), clientRequest.Header)
	if err != nil {
		log.Println(err, resp)
		w.WriteHeader(502)
		w.Write([]byte{})
		return
	}

	for k := range hopHeaders {
		delete(resp.Header, hopHeaders[k])
	}

	client_conn, err := websocket.Upgrade(w, clientRequest, resp.Header, 4096, 4096)
	if err != nil {
		log.Println(err)
		w.WriteHeader(502)
		w.Write([]byte{})
		return
	}

	go func() {
		for {
			messageType, p, err := client_conn.ReadMessage()
			if err != nil {
				log.Println("error while reading from client:", err)
				break
			}
			if err = upstream_conn.WriteMessage(messageType, p); err != nil {
				log.Println("error while writing to upstream:", err)
				break
			}
		}
		upstream_conn.Close()
	}()

	go func() {
		for {
			messageType, p, err := upstream_conn.ReadMessage()
			if err != nil {
				log.Println("error while reading from upstream:", err)
				break
			}
			if err = client_conn.WriteMessage(messageType, p); err != nil {
				log.Println("error while writing to client:", err)
				break
			}
		}

		client_conn.Close()
	}()
}

func writeResponseHeader(w http.ResponseWriter, r *http.Response) {
	for k := range r.Header {
		should_drop := false
		for i := range hopHeaders {
			if k == hopHeaders[i] {
				should_drop = true
				break
			}
		}

		if !should_drop {
			w.Header()[k] = r.Header[k]
		}
	}

	w.Header().Set("X-Forwarded-For", "127.0.0.1")
	w.WriteHeader(r.StatusCode)
}
<|endoftext|>


<commit_before>/* Copyright 2015 LinkedIn Corp. 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.
 */

package main

import (
	"bytes"
	log "github.com/cihub/seelog"
	"github.com/pborman/uuid"
	"net"
	"net/http"
	"os"
	"strings"
	"text/template"
	"time"
)

type HttpNotifier struct {
	app            *ApplicationContext
	templatePost   *template.Template
	templateDelete *template.Template
	extras         map[string]string
	ticker         *time.Ticker
	quitChan       chan struct{}
	groupIds       map[string]map[string]Event
	resultsChannel chan *ConsumerGroupStatus
	httpClient     *http.Client
}

type Event struct {
	Id    string
	Start time.Time
}

func NewHttpNotifier(app *ApplicationContext) (*HttpNotifier, error) {
	// Helper functions for templates
	fmap := template.FuncMap{
		"jsonencoder":    templateJsonEncoder,
		"topicsbystatus": classifyTopicsByStatus,
	}

	// Compile the templates
	templatePost, err := template.New("post").Funcs(fmap).ParseFiles(app.Config.Httpnotifier.TemplatePost)
	if err != nil {
		log.Criticalf("Cannot parse HTTP notifier POST template: %v", err)
		os.Exit(1)
	}
	templatePost = templatePost.Templates()[0]

	templateDelete, err := template.New("delete").Funcs(fmap).ParseFiles(app.Config.Httpnotifier.TemplateDelete)
	if err != nil {
		log.Criticalf("Cannot parse HTTP notifier DELETE template: %v", err)
		os.Exit(1)
	}
	templateDelete = templateDelete.Templates()[0]

	// Parse the extra parameters for the templates
	extras := make(map[string]string)
	for _, extra := range app.Config.Httpnotifier.Extras {
		parts := strings.Split(extra, "=")
		extras[parts[0]] = parts[1]
	}

	return &HttpNotifier{
		app:            app,
		templatePost:   templatePost,
		templateDelete: templateDelete,
		extras:         extras,
		quitChan:       make(chan struct{}),
		groupIds:       make(map[string]map[string]Event),
		resultsChannel: make(chan *ConsumerGroupStatus),
		httpClient: &http.Client{
			Transport: &http.Transport{
				Dial: (&net.Dialer{
					Timeout:   time.Duration(app.Config.Httpnotifier.Timeout) * time.Second,
					KeepAlive: time.Duration(app.Config.Httpnotifier.Keepalive) * time.Second,
				}).Dial,
				TLSHandshakeTimeout: time.Duration(app.Config.Httpnotifier.Timeout) * time.Second,
			},
		},
	}, nil
}

func (notifier *HttpNotifier) sendEvaluationRequests() {
	for cluster, _ := range notifier.app.Config.Kafka {
		// Get a current list of consumer groups
		storageRequest := &RequestConsumerList{Result: make(chan []string), Cluster: cluster}
		notifier.app.Storage.requestChannel <- storageRequest
		groups := <-storageRequest.Result

		// Send requests for group status
		for _, group := range groups {
			storageRequest := &RequestConsumerStatus{Result: notifier.resultsChannel, Cluster: cluster, Group: group}
			notifier.app.Storage.requestChannel <- storageRequest
		}
	}
}

func (notifier *HttpNotifier) handleEvaluationResponse(result *ConsumerGroupStatus) {
	if result.Status >= StatusWarning {
		if _, ok := notifier.groupIds[result.Cluster]; !ok {
			// Create the cluster map
			notifier.groupIds[result.Cluster] = make(map[string]Event)
		}
		if _, ok := notifier.groupIds[result.Cluster][result.Group]; !ok {
			// Create Event and Id
			eventId := uuid.NewRandom()
			notifier.groupIds[result.Cluster][result.Group] = Event{
				Id:    eventId.String(),
				Start: time.Now(),
			}
		}

		// NOTE - I'm leaving the JsonEncode item in here so as not to break compatibility. New helpers go in the FuncMap above
		bytesToSend := new(bytes.Buffer)
		err := notifier.templatePost.Execute(bytesToSend, struct {
			Cluster    string
			Group      string
			Id         string
			Start      time.Time
			Extras     map[string]string
			Result     *ConsumerGroupStatus
			JsonEncode func(interface{}) string
		}{
			Cluster:    result.Cluster,
			Group:      result.Group,
			Id:         notifier.groupIds[result.Cluster][result.Group].Id,
			Start:      notifier.groupIds[result.Cluster][result.Group].Start,
			Extras:     notifier.extras,
			Result:     result,
			JsonEncode: templateJsonEncoder,
		})
		if err != nil {
			log.Errorf("Failed to assemble POST: %v", err)
			return
		}

		// Send POST to HTTP endpoint
		req, err := http.NewRequest("POST", notifier.app.Config.Httpnotifier.Url, bytesToSend)
		req.Header.Set("Content-Type", "application/json")

		resp, err := notifier.httpClient.Do(req)
		if err != nil {
			log.Errorf("Failed to send POST (Id %s): %v", notifier.groupIds[result.Cluster][result.Group].Id, err)
			return
		}
		defer resp.Body.Close()

		if (resp.StatusCode >= 200) && (resp.StatusCode <= 299) {
			log.Infof("Sent POST for group %s in cluster %s at severity %v (Id %s)", result.Group,
				result.Cluster, result.Status, notifier.groupIds[result.Cluster][result.Group].Id)
		} else {
			log.Errorf("Failed to send POST for group %s in cluster %s at severity %v (Id %s): %s", result.Group,
				result.Cluster, result.Status, notifier.groupIds[result.Cluster][result.Group].Id, resp.Status)
		}
	} else {
		if _, ok := notifier.groupIds[result.Cluster][result.Group]; ok {
			// Send DELETE to HTTP endpoint
			bytesToSend := new(bytes.Buffer)
			err := notifier.templateDelete.Execute(bytesToSend, struct {
				Cluster string
				Group   string
				Id      string
				Start   time.Time
				Extras  map[string]string
			}{
				Cluster: result.Cluster,
				Group:   result.Group,
				Id:      notifier.groupIds[result.Cluster][result.Group].Id,
				Start:   notifier.groupIds[result.Cluster][result.Group].Start,
				Extras:  notifier.extras,
			})
			if err != nil {
				log.Errorf("Failed to assemble DELETE: %v", err)
				return
			}

			req, err := http.NewRequest("DELETE", notifier.app.Config.Httpnotifier.Url, bytesToSend)
			req.Header.Set("Content-Type", "application/json")

			resp, err := notifier.httpClient.Do(req)
			if err != nil {
				log.Errorf("Failed to send DELETE: %v", err)
				return
			}
			defer resp.Body.Close()

			if (resp.StatusCode >= 200) && (resp.StatusCode <= 299) {
				log.Infof("Sent DELETE for group %s in cluster %s (Id %s)", result.Group, result.Cluster,
					notifier.groupIds[result.Cluster][result.Group].Id)
			} else {
				log.Errorf("Failed to send DELETE for group %s in cluster %s (Id %s): %s", result.Group,
					result.Cluster, notifier.groupIds[result.Cluster][result.Group].Id, resp.Status)
			}

			// Remove ID for group that is now clear
			delete(notifier.groupIds[result.Cluster], result.Group)
		}
	}
}

func (notifier *HttpNotifier) Start() {
	// Start the ticker
	notifier.ticker = time.NewTicker(time.Duration(notifier.app.Config.Httpnotifier.Interval) * time.Second)

	go func() {
	OUTERLOOP:
		for {
			select {
			case <-notifier.quitChan:
				notifier.ticker.Stop()
				break OUTERLOOP
			case <-notifier.ticker.C:
				go notifier.sendEvaluationRequests()
			case result := <-notifier.resultsChannel:
				go notifier.handleEvaluationResponse(result)
			}
		}
	}()
}

func (notifier *HttpNotifier) Stop() {
	close(notifier.quitChan)
}
<commit_msg>Always read response bodies before closing. Clean up timeout in the client<commit_after>/* Copyright 2015 LinkedIn Corp. 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.
 */

package main

import (
	"bytes"
	log "github.com/cihub/seelog"
	"github.com/pborman/uuid"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strings"
	"text/template"
	"time"
)

type HttpNotifier struct {
	app            *ApplicationContext
	templatePost   *template.Template
	templateDelete *template.Template
	extras         map[string]string
	ticker         *time.Ticker
	quitChan       chan struct{}
	groupIds       map[string]map[string]Event
	resultsChannel chan *ConsumerGroupStatus
	httpClient     *http.Client
}

type Event struct {
	Id    string
	Start time.Time
}

func NewHttpNotifier(app *ApplicationContext) (*HttpNotifier, error) {
	// Helper functions for templates
	fmap := template.FuncMap{
		"jsonencoder":    templateJsonEncoder,
		"topicsbystatus": classifyTopicsByStatus,
	}

	// Compile the templates
	templatePost, err := template.New("post").Funcs(fmap).ParseFiles(app.Config.Httpnotifier.TemplatePost)
	if err != nil {
		log.Criticalf("Cannot parse HTTP notifier POST template: %v", err)
		os.Exit(1)
	}
	templatePost = templatePost.Templates()[0]

	templateDelete, err := template.New("delete").Funcs(fmap).ParseFiles(app.Config.Httpnotifier.TemplateDelete)
	if err != nil {
		log.Criticalf("Cannot parse HTTP notifier DELETE template: %v", err)
		os.Exit(1)
	}
	templateDelete = templateDelete.Templates()[0]

	// Parse the extra parameters for the templates
	extras := make(map[string]string)
	for _, extra := range app.Config.Httpnotifier.Extras {
		parts := strings.Split(extra, "=")
		extras[parts[0]] = parts[1]
	}

	return &HttpNotifier{
		app:            app,
		templatePost:   templatePost,
		templateDelete: templateDelete,
		extras:         extras,
		quitChan:       make(chan struct{}),
		groupIds:       make(map[string]map[string]Event),
		resultsChannel: make(chan *ConsumerGroupStatus),
		httpClient: &http.Client{
			Timeout: time.Duration(app.Config.Httpnotifier.Timeout) * time.Second,
			Transport: &http.Transport{
				Dial: (&net.Dialer{
					KeepAlive: time.Duration(app.Config.Httpnotifier.Keepalive) * time.Second,
				}).Dial,
			},
		},
	}, nil
}

func (notifier *HttpNotifier) sendEvaluationRequests() {
	for cluster, _ := range notifier.app.Config.Kafka {
		// Get a current list of consumer groups
		storageRequest := &RequestConsumerList{Result: make(chan []string), Cluster: cluster}
		notifier.app.Storage.requestChannel <- storageRequest
		groups := <-storageRequest.Result

		// Send requests for group status
		for _, group := range groups {
			storageRequest := &RequestConsumerStatus{Result: notifier.resultsChannel, Cluster: cluster, Group: group}
			notifier.app.Storage.requestChannel <- storageRequest
		}
	}
}

func (notifier *HttpNotifier) handleEvaluationResponse(result *ConsumerGroupStatus) {
	if result.Status >= StatusWarning {
		if _, ok := notifier.groupIds[result.Cluster]; !ok {
			// Create the cluster map
			notifier.groupIds[result.Cluster] = make(map[string]Event)
		}
		if _, ok := notifier.groupIds[result.Cluster][result.Group]; !ok {
			// Create Event and Id
			eventId := uuid.NewRandom()
			notifier.groupIds[result.Cluster][result.Group] = Event{
				Id:    eventId.String(),
				Start: time.Now(),
			}
		}

		// NOTE - I'm leaving the JsonEncode item in here so as not to break compatibility. New helpers go in the FuncMap above
		bytesToSend := new(bytes.Buffer)
		err := notifier.templatePost.Execute(bytesToSend, struct {
			Cluster    string
			Group      string
			Id         string
			Start      time.Time
			Extras     map[string]string
			Result     *ConsumerGroupStatus
			JsonEncode func(interface{}) string
		}{
			Cluster:    result.Cluster,
			Group:      result.Group,
			Id:         notifier.groupIds[result.Cluster][result.Group].Id,
			Start:      notifier.groupIds[result.Cluster][result.Group].Start,
			Extras:     notifier.extras,
			Result:     result,
			JsonEncode: templateJsonEncoder,
		})
		if err != nil {
			log.Errorf("Failed to assemble POST: %v", err)
			return
		}

		// Send POST to HTTP endpoint
		req, err := http.NewRequest("POST", notifier.app.Config.Httpnotifier.Url, bytesToSend)
		req.Header.Set("Content-Type", "application/json")

		resp, err := notifier.httpClient.Do(req)
		if err != nil {
			log.Errorf("Failed to send POST (Id %s): %v", notifier.groupIds[result.Cluster][result.Group].Id, err)
			return
		}
		io.Copy(ioutil.Discard, resp.Body)
		resp.Body.Close()

		if (resp.StatusCode >= 200) && (resp.StatusCode <= 299) {
			log.Infof("Sent POST for group %s in cluster %s at severity %v (Id %s)", result.Group,
				result.Cluster, result.Status, notifier.groupIds[result.Cluster][result.Group].Id)
		} else {
			log.Errorf("Failed to send POST for group %s in cluster %s at severity %v (Id %s): %s", result.Group,
				result.Cluster, result.Status, notifier.groupIds[result.Cluster][result.Group].Id, resp.Status)
		}
	} else {
		if _, ok := notifier.groupIds[result.Cluster][result.Group]; ok {
			// Send DELETE to HTTP endpoint
			bytesToSend := new(bytes.Buffer)
			err := notifier.templateDelete.Execute(bytesToSend, struct {
				Cluster string
				Group   string
				Id      string
				Start   time.Time
				Extras  map[string]string
			}{
				Cluster: result.Cluster,
				Group:   result.Group,
				Id:      notifier.groupIds[result.Cluster][result.Group].Id,
				Start:   notifier.groupIds[result.Cluster][result.Group].Start,
				Extras:  notifier.extras,
			})
			if err != nil {
				log.Errorf("Failed to assemble DELETE: %v", err)
				return
			}

			req, err := http.NewRequest("DELETE", notifier.app.Config.Httpnotifier.Url, bytesToSend)
			req.Header.Set("Content-Type", "application/json")

			resp, err := notifier.httpClient.Do(req)
			if err != nil {
				log.Errorf("Failed to send DELETE: %v", err)
				return
			}
			io.Copy(ioutil.Discard, resp.Body)
			resp.Body.Close()

			if (resp.StatusCode >= 200) && (resp.StatusCode <= 299) {
				log.Infof("Sent DELETE for group %s in cluster %s (Id %s)", result.Group, result.Cluster,
					notifier.groupIds[result.Cluster][result.Group].Id)
			} else {
				log.Errorf("Failed to send DELETE for group %s in cluster %s (Id %s): %s", result.Group,
					result.Cluster, notifier.groupIds[result.Cluster][result.Group].Id, resp.Status)
			}

			// Remove ID for group that is now clear
			delete(notifier.groupIds[result.Cluster], result.Group)
		}
	}
}

func (notifier *HttpNotifier) Start() {
	// Start the ticker
	notifier.ticker = time.NewTicker(time.Duration(notifier.app.Config.Httpnotifier.Interval) * time.Second)

	go func() {
	OUTERLOOP:
		for {
			select {
			case <-notifier.quitChan:
				notifier.ticker.Stop()
				break OUTERLOOP
			case <-notifier.ticker.C:
				go notifier.sendEvaluationRequests()
			case result := <-notifier.resultsChannel:
				go notifier.handleEvaluationResponse(result)
			}
		}
	}()
}

func (notifier *HttpNotifier) Stop() {
	close(notifier.quitChan)
}
<|endoftext|>


<commit_before>package docker

import (
	"github.com/dotcloud/docker"
	"github.com/dotcloud/docker/utils"
	"io/ioutil"
	"strings"
	"testing"
)

func TestContainerTagImageDelete(t *testing.T) {
	eng := NewTestEngine(t)
	defer mkRuntimeFromEngine(eng, t).Nuke()

	srv := mkServerFromEngine(eng, t)

	initialImages, err := srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}
	if err := srv.ContainerTag(unitTestImageName, "utest", "tag1", false); err != nil {
		t.Fatal(err)
	}

	if err := srv.ContainerTag(unitTestImageName, "utest/docker", "tag2", false); err != nil {
		t.Fatal(err)
	}
	if err := srv.ContainerTag(unitTestImageName, "utest:5000/docker", "tag3", false); err != nil {
		t.Fatal(err)
	}

	images, err := srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}

	if len(images[0].RepoTags) != len(initialImages[0].RepoTags)+3 {
		t.Errorf("Expected %d images, %d found", len(initialImages)+3, len(images))
	}

	if _, err := srv.ImageDelete("utest/docker:tag2", true); err != nil {
		t.Fatal(err)
	}

	images, err = srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}

	if len(images[0].RepoTags) != len(initialImages[0].RepoTags)+2 {
		t.Errorf("Expected %d images, %d found", len(initialImages)+2, len(images))
	}

	if _, err := srv.ImageDelete("utest:5000/docker:tag3", true); err != nil {
		t.Fatal(err)
	}

	images, err = srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}

	if len(images[0].RepoTags) != len(initialImages[0].RepoTags)+1 {
		t.Errorf("Expected %d images, %d found", len(initialImages)+1, len(images))
	}

	if _, err := srv.ImageDelete("utest:tag1", true); err != nil {
		t.Fatal(err)
	}

	images, err = srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}

	if len(images) != len(initialImages) {
		t.Errorf("Expected %d image, %d found", len(initialImages), len(images))
	}
}

func TestCreateRm(t *testing.T) {
	eng := NewTestEngine(t)
	srv := mkServerFromEngine(eng, t)
	defer mkRuntimeFromEngine(eng, t).Nuke()

	config, _, _, err := docker.ParseRun([]string{unitTestImageID, "echo test"}, nil)
	if err != nil {
		t.Fatal(err)
	}

	id := createTestContainer(eng, config, t)

	if c := srv.Containers(true, false, -1, "", ""); len(c) != 1 {
		t.Errorf("Expected 1 container, %v found", len(c))
	}

	if err = srv.ContainerDestroy(id, true, false); err != nil {
		t.Fatal(err)
	}

	if c := srv.Containers(true, false, -1, "", ""); len(c) != 0 {
		t.Errorf("Expected 0 container, %v found", len(c))
	}

}

func TestCreateRmVolumes(t *testing.T) {
	eng := NewTestEngine(t)
	srv := mkServerFromEngine(eng, t)
	defer mkRuntimeFromEngine(eng, t).Nuke()

	config, hostConfig, _, err := docker.ParseRun([]string{"-v", "/srv", unitTestImageID, "echo", "test"}, nil)
	if err != nil {
		t.Fatal(err)
	}

	id := createTestContainer(eng, config, t)

	if c := srv.Containers(true, false, -1, "", ""); len(c) != 1 {
		t.Errorf("Expected 1 container, %v found", len(c))
	}

	job := eng.Job("start", id)
	if err := job.ImportEnv(hostConfig); err != nil {
		t.Fatal(err)
	}
	if err := job.Run(); err != nil {
		t.Fatal(err)
	}

	err = srv.ContainerStop(id, 1)
	if err != nil {
		t.Fatal(err)
	}

	if err = srv.ContainerDestroy(id, true, false); err != nil {
		t.Fatal(err)
	}

	if c := srv.Containers(true, false, -1, "", ""); len(c) != 0 {
		t.Errorf("Expected 0 container, %v found", len(c))
	}
}

func TestCommit(t *testing.T) {
	eng := NewTestEngine(t)
	srv := mkServerFromEngine(eng, t)
	defer mkRuntimeFromEngine(eng, t).Nuke()

	config, _, _, err := docker.ParseRun([]string{unitTestImageID, "/bin/cat"}, nil)
	if err != nil {
		t.Fatal(err)
	}

	id := createTestContainer(eng, config, t)

	if _, err := srv.ContainerCommit(id, "testrepo", "testtag", "", "", config); err != nil {
		t.Fatal(err)
	}
}

func TestCreateStartRestartStopStartKillRm(t *testing.T) {
	eng := NewTestEngine(t)
	srv := mkServerFromEngine(eng, t)
	defer mkRuntimeFromEngine(eng, t).Nuke()

	config, hostConfig, _, err := docker.ParseRun([]string{"-i", unitTestImageID, "/bin/cat"}, nil)
	if err != nil {
		t.Fatal(err)
	}

	id := createTestContainer(eng, config, t)

	if c := srv.Containers(true, false, -1, "", ""); len(c) != 1 {
		t.Errorf("Expected 1 container, %v found", len(c))
	}

	job := eng.Job("start", id)
	if err := job.ImportEnv(hostConfig); err != nil {
		t.Fatal(err)
	}
	if err := job.Run(); err != nil {
		t.Fatal(err)
	}

	if err := srv.ContainerRestart(id, 15); err != nil {
		t.Fatal(err)
	}

	if err := srv.ContainerStop(id, 15); err != nil {
		t.Fatal(err)
	}

	job = eng.Job("start", id)
	if err := job.ImportEnv(hostConfig); err != nil {
		t.Fatal(err)
	}
	if err := job.Run(); err != nil {
		t.Fatal(err)
	}

	if err := srv.ContainerKill(id, 0); err != nil {
		t.Fatal(err)
	}

	// FIXME: this failed once with a race condition ("Unable to remove filesystem for xxx: directory not empty")
	if err := srv.ContainerDestroy(id, true, false); err != nil {
		t.Fatal(err)
	}

	if c := srv.Containers(true, false, -1, "", ""); len(c) != 0 {
		t.Errorf("Expected 0 container, %v found", len(c))
	}
}

func TestRunWithTooLowMemoryLimit(t *testing.T) {
	eng := NewTestEngine(t)
	defer mkRuntimeFromEngine(eng, t).Nuke()

	// Try to create a container with a memory limit of 1 byte less than the minimum allowed limit.
	job := eng.Job("create")
	job.Setenv("Image", unitTestImageID)
	job.Setenv("Memory", "524287")
	job.Setenv("CpuShares", "1000")
	job.SetenvList("Cmd", []string{"/bin/cat"})
	var id string
	job.StdoutParseString(&id)
	if err := job.Run(); err == nil {
		t.Errorf("Memory limit is smaller than the allowed limit. Container creation should've failed!")
	}
}

func TestRmi(t *testing.T) {
	eng := NewTestEngine(t)
	srv := mkServerFromEngine(eng, t)
	defer mkRuntimeFromEngine(eng, t).Nuke()

	initialImages, err := srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}

	config, hostConfig, _, err := docker.ParseRun([]string{unitTestImageID, "echo", "test"}, nil)
	if err != nil {
		t.Fatal(err)
	}

	containerID := createTestContainer(eng, config, t)

	//To remove
	job := eng.Job("start", containerID)
	if err := job.ImportEnv(hostConfig); err != nil {
		t.Fatal(err)
	}
	if err := job.Run(); err != nil {
		t.Fatal(err)
	}

	imageID, err := srv.ContainerCommit(containerID, "test", "", "", "", nil)
	if err != nil {
		t.Fatal(err)
	}

	err = srv.ContainerTag(imageID, "test", "0.1", false)
	if err != nil {
		t.Fatal(err)
	}

	containerID = createTestContainer(eng, config, t)

	//To remove
	job = eng.Job("start", containerID)
	if err := job.ImportEnv(hostConfig); err != nil {
		t.Fatal(err)
	}
	if err := job.Run(); err != nil {
		t.Fatal(err)
	}

	_, err = srv.ContainerCommit(containerID, "test", "", "", "", nil)
	if err != nil {
		t.Fatal(err)
	}

	images, err := srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}

	if len(images)-len(initialImages) != 2 {
		t.Fatalf("Expected 2 new images, found %d.", len(images)-len(initialImages))
	}

	_, err = srv.ImageDelete(imageID, true)
	if err != nil {
		t.Fatal(err)
	}

	images, err = srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}

	if len(images)-len(initialImages) != 1 {
		t.Fatalf("Expected 1 new image, found %d.", len(images)-len(initialImages))
	}

	for _, image := range images {
		if strings.Contains(unitTestImageID, image.ID) {
			continue
		}
		if image.RepoTags[0] == "<none>:<none>" {
			t.Fatalf("Expected tagged image, got untagged one.")
		}
	}
}

func TestImagesFilter(t *testing.T) {
	eng := NewTestEngine(t)
	defer nuke(mkRuntimeFromEngine(eng, t))

	srv := mkServerFromEngine(eng, t)

	if err := srv.ContainerTag(unitTestImageName, "utest", "tag1", false); err != nil {
		t.Fatal(err)
	}

	if err := srv.ContainerTag(unitTestImageName, "utest/docker", "tag2", false); err != nil {
		t.Fatal(err)
	}
	if err := srv.ContainerTag(unitTestImageName, "utest:5000/docker", "tag3", false); err != nil {
		t.Fatal(err)
	}

	images, err := srv.Images(false, "utest*/*")
	if err != nil {
		t.Fatal(err)
	}

	if len(images[0].RepoTags) != 2 {
		t.Fatal("incorrect number of matches returned")
	}

	images, err = srv.Images(false, "utest")
	if err != nil {
		t.Fatal(err)
	}

	if len(images[0].RepoTags) != 1 {
		t.Fatal("incorrect number of matches returned")
	}

	images, err = srv.Images(false, "utest*")
	if err != nil {
		t.Fatal(err)
	}

	if len(images[0].RepoTags) != 1 {
		t.Fatal("incorrect number of matches returned")
	}

	images, err = srv.Images(false, "*5000*/*")
	if err != nil {
		t.Fatal(err)
	}

	if len(images[0].RepoTags) != 1 {
		t.Fatal("incorrect number of matches returned")
	}
}

func TestImageInsert(t *testing.T) {
	eng := NewTestEngine(t)
	defer mkRuntimeFromEngine(eng, t).Nuke()
	srv := mkServerFromEngine(eng, t)
	sf := utils.NewStreamFormatter(true)

	// bad image name fails
	if err := srv.ImageInsert("foo", "https://www.docker.io/static/img/docker-top-logo.png", "/foo", ioutil.Discard, sf); err == nil {
		t.Fatal("expected an error and got none")
	}

	// bad url fails
	if err := srv.ImageInsert(unitTestImageID, "http://bad_host_name_that_will_totally_fail.com/", "/foo", ioutil.Discard, sf); err == nil {
		t.Fatal("expected an error and got none")
	}

	// success returns nil
	if err := srv.ImageInsert(unitTestImageID, "https://www.docker.io/static/img/docker-top-logo.png", "/foo", ioutil.Discard, sf); err != nil {
		t.Fatalf("expected no error, but got %v", err)
	}
}
<commit_msg>test: fix TestRmi race condition<commit_after>package docker

import (
	"github.com/dotcloud/docker"
	"github.com/dotcloud/docker/utils"
	"io/ioutil"
	"strings"
	"testing"
)

func TestContainerTagImageDelete(t *testing.T) {
	eng := NewTestEngine(t)
	defer mkRuntimeFromEngine(eng, t).Nuke()

	srv := mkServerFromEngine(eng, t)

	initialImages, err := srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}
	if err := srv.ContainerTag(unitTestImageName, "utest", "tag1", false); err != nil {
		t.Fatal(err)
	}

	if err := srv.ContainerTag(unitTestImageName, "utest/docker", "tag2", false); err != nil {
		t.Fatal(err)
	}
	if err := srv.ContainerTag(unitTestImageName, "utest:5000/docker", "tag3", false); err != nil {
		t.Fatal(err)
	}

	images, err := srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}

	if len(images[0].RepoTags) != len(initialImages[0].RepoTags)+3 {
		t.Errorf("Expected %d images, %d found", len(initialImages)+3, len(images))
	}

	if _, err := srv.ImageDelete("utest/docker:tag2", true); err != nil {
		t.Fatal(err)
	}

	images, err = srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}

	if len(images[0].RepoTags) != len(initialImages[0].RepoTags)+2 {
		t.Errorf("Expected %d images, %d found", len(initialImages)+2, len(images))
	}

	if _, err := srv.ImageDelete("utest:5000/docker:tag3", true); err != nil {
		t.Fatal(err)
	}

	images, err = srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}

	if len(images[0].RepoTags) != len(initialImages[0].RepoTags)+1 {
		t.Errorf("Expected %d images, %d found", len(initialImages)+1, len(images))
	}

	if _, err := srv.ImageDelete("utest:tag1", true); err != nil {
		t.Fatal(err)
	}

	images, err = srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}

	if len(images) != len(initialImages) {
		t.Errorf("Expected %d image, %d found", len(initialImages), len(images))
	}
}

func TestCreateRm(t *testing.T) {
	eng := NewTestEngine(t)
	srv := mkServerFromEngine(eng, t)
	defer mkRuntimeFromEngine(eng, t).Nuke()

	config, _, _, err := docker.ParseRun([]string{unitTestImageID, "echo test"}, nil)
	if err != nil {
		t.Fatal(err)
	}

	id := createTestContainer(eng, config, t)

	if c := srv.Containers(true, false, -1, "", ""); len(c) != 1 {
		t.Errorf("Expected 1 container, %v found", len(c))
	}

	if err = srv.ContainerDestroy(id, true, false); err != nil {
		t.Fatal(err)
	}

	if c := srv.Containers(true, false, -1, "", ""); len(c) != 0 {
		t.Errorf("Expected 0 container, %v found", len(c))
	}

}

func TestCreateRmVolumes(t *testing.T) {
	eng := NewTestEngine(t)
	srv := mkServerFromEngine(eng, t)
	defer mkRuntimeFromEngine(eng, t).Nuke()

	config, hostConfig, _, err := docker.ParseRun([]string{"-v", "/srv", unitTestImageID, "echo", "test"}, nil)
	if err != nil {
		t.Fatal(err)
	}

	id := createTestContainer(eng, config, t)

	if c := srv.Containers(true, false, -1, "", ""); len(c) != 1 {
		t.Errorf("Expected 1 container, %v found", len(c))
	}

	job := eng.Job("start", id)
	if err := job.ImportEnv(hostConfig); err != nil {
		t.Fatal(err)
	}
	if err := job.Run(); err != nil {
		t.Fatal(err)
	}

	err = srv.ContainerStop(id, 1)
	if err != nil {
		t.Fatal(err)
	}

	if err = srv.ContainerDestroy(id, true, false); err != nil {
		t.Fatal(err)
	}

	if c := srv.Containers(true, false, -1, "", ""); len(c) != 0 {
		t.Errorf("Expected 0 container, %v found", len(c))
	}
}

func TestCommit(t *testing.T) {
	eng := NewTestEngine(t)
	srv := mkServerFromEngine(eng, t)
	defer mkRuntimeFromEngine(eng, t).Nuke()

	config, _, _, err := docker.ParseRun([]string{unitTestImageID, "/bin/cat"}, nil)
	if err != nil {
		t.Fatal(err)
	}

	id := createTestContainer(eng, config, t)

	if _, err := srv.ContainerCommit(id, "testrepo", "testtag", "", "", config); err != nil {
		t.Fatal(err)
	}
}

func TestCreateStartRestartStopStartKillRm(t *testing.T) {
	eng := NewTestEngine(t)
	srv := mkServerFromEngine(eng, t)
	defer mkRuntimeFromEngine(eng, t).Nuke()

	config, hostConfig, _, err := docker.ParseRun([]string{"-i", unitTestImageID, "/bin/cat"}, nil)
	if err != nil {
		t.Fatal(err)
	}

	id := createTestContainer(eng, config, t)

	if c := srv.Containers(true, false, -1, "", ""); len(c) != 1 {
		t.Errorf("Expected 1 container, %v found", len(c))
	}

	job := eng.Job("start", id)
	if err := job.ImportEnv(hostConfig); err != nil {
		t.Fatal(err)
	}
	if err := job.Run(); err != nil {
		t.Fatal(err)
	}

	if err := srv.ContainerRestart(id, 15); err != nil {
		t.Fatal(err)
	}

	if err := srv.ContainerStop(id, 15); err != nil {
		t.Fatal(err)
	}

	job = eng.Job("start", id)
	if err := job.ImportEnv(hostConfig); err != nil {
		t.Fatal(err)
	}
	if err := job.Run(); err != nil {
		t.Fatal(err)
	}

	if err := srv.ContainerKill(id, 0); err != nil {
		t.Fatal(err)
	}

	// FIXME: this failed once with a race condition ("Unable to remove filesystem for xxx: directory not empty")
	if err := srv.ContainerDestroy(id, true, false); err != nil {
		t.Fatal(err)
	}

	if c := srv.Containers(true, false, -1, "", ""); len(c) != 0 {
		t.Errorf("Expected 0 container, %v found", len(c))
	}
}

func TestRunWithTooLowMemoryLimit(t *testing.T) {
	eng := NewTestEngine(t)
	defer mkRuntimeFromEngine(eng, t).Nuke()

	// Try to create a container with a memory limit of 1 byte less than the minimum allowed limit.
	job := eng.Job("create")
	job.Setenv("Image", unitTestImageID)
	job.Setenv("Memory", "524287")
	job.Setenv("CpuShares", "1000")
	job.SetenvList("Cmd", []string{"/bin/cat"})
	var id string
	job.StdoutParseString(&id)
	if err := job.Run(); err == nil {
		t.Errorf("Memory limit is smaller than the allowed limit. Container creation should've failed!")
	}
}

func TestRmi(t *testing.T) {
	eng := NewTestEngine(t)
	srv := mkServerFromEngine(eng, t)
	defer mkRuntimeFromEngine(eng, t).Nuke()

	initialImages, err := srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}

	config, hostConfig, _, err := docker.ParseRun([]string{unitTestImageID, "echo", "test"}, nil)
	if err != nil {
		t.Fatal(err)
	}

	containerID := createTestContainer(eng, config, t)

	//To remove
	job := eng.Job("start", containerID)
	if err := job.ImportEnv(hostConfig); err != nil {
		t.Fatal(err)
	}
	if err := job.Run(); err != nil {
		t.Fatal(err)
	}

	if _, err := srv.ContainerWait(containerID); err != nil {
		t.Fatal(err)
	}

	imageID, err := srv.ContainerCommit(containerID, "test", "", "", "", nil)
	if err != nil {
		t.Fatal(err)
	}

	err = srv.ContainerTag(imageID, "test", "0.1", false)
	if err != nil {
		t.Fatal(err)
	}

	containerID = createTestContainer(eng, config, t)

	//To remove
	job = eng.Job("start", containerID)
	if err := job.ImportEnv(hostConfig); err != nil {
		t.Fatal(err)
	}
	if err := job.Run(); err != nil {
		t.Fatal(err)
	}

	if _, err := srv.ContainerWait(containerID); err != nil {
		t.Fatal(err)
	}

	_, err = srv.ContainerCommit(containerID, "test", "", "", "", nil)
	if err != nil {
		t.Fatal(err)
	}

	images, err := srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}

	if len(images)-len(initialImages) != 2 {
		t.Fatalf("Expected 2 new images, found %d.", len(images)-len(initialImages))
	}

	_, err = srv.ImageDelete(imageID, true)
	if err != nil {
		t.Fatal(err)
	}

	images, err = srv.Images(false, "")
	if err != nil {
		t.Fatal(err)
	}

	if len(images)-len(initialImages) != 1 {
		t.Fatalf("Expected 1 new image, found %d.", len(images)-len(initialImages))
	}

	for _, image := range images {
		if strings.Contains(unitTestImageID, image.ID) {
			continue
		}
		if image.RepoTags[0] == "<none>:<none>" {
			t.Fatalf("Expected tagged image, got untagged one.")
		}
	}
}

func TestImagesFilter(t *testing.T) {
	eng := NewTestEngine(t)
	defer nuke(mkRuntimeFromEngine(eng, t))

	srv := mkServerFromEngine(eng, t)

	if err := srv.ContainerTag(unitTestImageName, "utest", "tag1", false); err != nil {
		t.Fatal(err)
	}

	if err := srv.ContainerTag(unitTestImageName, "utest/docker", "tag2", false); err != nil {
		t.Fatal(err)
	}
	if err := srv.ContainerTag(unitTestImageName, "utest:5000/docker", "tag3", false); err != nil {
		t.Fatal(err)
	}

	images, err := srv.Images(false, "utest*/*")
	if err != nil {
		t.Fatal(err)
	}

	if len(images[0].RepoTags) != 2 {
		t.Fatal("incorrect number of matches returned")
	}

	images, err = srv.Images(false, "utest")
	if err != nil {
		t.Fatal(err)
	}

	if len(images[0].RepoTags) != 1 {
		t.Fatal("incorrect number of matches returned")
	}

	images, err = srv.Images(false, "utest*")
	if err != nil {
		t.Fatal(err)
	}

	if len(images[0].RepoTags) != 1 {
		t.Fatal("incorrect number of matches returned")
	}

	images, err = srv.Images(false, "*5000*/*")
	if err != nil {
		t.Fatal(err)
	}

	if len(images[0].RepoTags) != 1 {
		t.Fatal("incorrect number of matches returned")
	}
}

func TestImageInsert(t *testing.T) {
	eng := NewTestEngine(t)
	defer mkRuntimeFromEngine(eng, t).Nuke()
	srv := mkServerFromEngine(eng, t)
	sf := utils.NewStreamFormatter(true)

	// bad image name fails
	if err := srv.ImageInsert("foo", "https://www.docker.io/static/img/docker-top-logo.png", "/foo", ioutil.Discard, sf); err == nil {
		t.Fatal("expected an error and got none")
	}

	// bad url fails
	if err := srv.ImageInsert(unitTestImageID, "http://bad_host_name_that_will_totally_fail.com/", "/foo", ioutil.Discard, sf); err == nil {
		t.Fatal("expected an error and got none")
	}

	// success returns nil
	if err := srv.ImageInsert(unitTestImageID, "https://www.docker.io/static/img/docker-top-logo.png", "/foo", ioutil.Discard, sf); err != nil {
		t.Fatalf("expected no error, but got %v", err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 CoreOS, Inc.
//
// 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 util

import (
	"encoding/hex"
	"hash"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"path/filepath"
	"strconv"

	"github.com/coreos/ignition/config/types"
	"github.com/coreos/ignition/internal/log"
	"github.com/coreos/ignition/internal/resource"
)

const (
	DefaultDirectoryPermissions os.FileMode = 0755
	DefaultFilePermissions      os.FileMode = 0644
)

type FetchOp struct {
	Hash         hash.Hash
	Path         string
	Mode         os.FileMode
	Uid          int
	Gid          int
	Url          url.URL
	FetchOptions resource.FetchOptions
}

// newHashedReader returns a new ReadCloser that also writes to the provided hash.
func newHashedReader(reader io.ReadCloser, hasher hash.Hash) io.ReadCloser {
	return struct {
		io.Reader
		io.Closer
	}{
		Reader: io.TeeReader(reader, hasher),
		Closer: reader,
	}
}

// PrepareFetch converts a given logger, http client, and types.File into a
// FetchOp. This includes operations such as parsing the source URL, generating
// a hasher, and performing user/group name lookups. If an error is encountered,
// the issue will be logged and nil will be returned.
func (u Util) PrepareFetch(l *log.Logger, f types.File) *FetchOp {
	var err error
	var expectedSum []byte

	// explicitly ignoring the error here because the config should already be
	// validated by this point
	uri, _ := url.Parse(f.Contents.Source)

	hasher, err := GetHasher(f.Contents.Verification)
	if err != nil {
		l.Crit("Error verifying file %q: %v", f.Path, err)
		return nil
	}

	if hasher != nil {
		// explicitly ignoring the error here because the config should already
		// be validated by this point
		_, expectedSumString, _ := f.Contents.Verification.HashParts()
		expectedSum, err = hex.DecodeString(expectedSumString)
		if err != nil {
			l.Crit("Error parsing verification string %q: %v", expectedSumString, err)
			return nil
		}
	}

	if f.User.Name != "" {
		user, err := u.userLookup(f.User.Name)
		if err != nil {
			l.Crit("No such user %q: %v", f.User.Name, err)
			return nil
		}
		uid, err := strconv.ParseInt(user.Uid, 0, 0)
		if err != nil {
			l.Crit("Couldn't parse uid %q: %v", user.Uid, err)
			return nil
		}
		tmp := int(uid)
		f.User.ID = &tmp
	}
	if f.Group.Name != "" {
		g, err := u.groupLookup(f.Group.Name)
		if err != nil {
			l.Crit("No such group %q: %v", f.Group.Name, err)
			return nil
		}
		gid, err := strconv.ParseInt(g.Gid, 0, 0)
		if err != nil {
			l.Crit("Couldn't parse gid %q: %v", g.Gid, err)
			return nil
		}
		tmp := int(gid)
		f.Group.ID = &tmp
	}

	return &FetchOp{
		Path: f.Path,
		Hash: hasher,
		Mode: os.FileMode(f.Mode),
		Uid:  *f.User.ID,
		Gid:  *f.Group.ID,
		Url:  *uri,
		FetchOptions: resource.FetchOptions{
			Hash:        hasher,
			Compression: f.Contents.Compression,
			ExpectedSum: expectedSum,
		},
	}
}

func (u Util) WriteLink(s types.Link) error {
	path := u.JoinPath(s.Path)

	if err := MkdirForFile(path); err != nil {
		return err
	}

	if s.Hard {
		targetPath := u.JoinPath(s.Target)
		return os.Link(targetPath, path)
	}
	return os.Symlink(s.Target, path)
}

// PerformFetch performs a fetch operation generated by PrepareFetch, retrieving
// the file and writing it to disk. Any encountered errors are returned.
func (u Util) PerformFetch(f *FetchOp) error {
	var err error

	path := u.JoinPath(string(f.Path))

	if err := MkdirForFile(path); err != nil {
		return err
	}

	// Create a temporary file in the same directory to ensure it's on the same filesystem
	var tmp *os.File
	if tmp, err = ioutil.TempFile(filepath.Dir(path), "tmp"); err != nil {
		return err
	}

	defer func() {
		tmp.Close()
		if err != nil {
			os.Remove(tmp.Name())
		}
	}()

	err = u.Fetcher.Fetch(f.Url, tmp, f.FetchOptions)
	if err != nil {
		u.Crit("Error fetching file %q: %v", f.Path, err)
		return err
	}

	// XXX(vc): Note that we assume to be operating on the file we just wrote, this is only guaranteed
	// by using syscall.Fchown() and syscall.Fchmod()

	// Ensure the ownership and mode are as requested (since WriteFile can be affected by sticky bit)
	if err = os.Chown(tmp.Name(), f.Uid, f.Gid); err != nil {
		return err
	}

	if err = os.Chmod(tmp.Name(), f.Mode); err != nil {
		return err
	}

	if err = os.Rename(tmp.Name(), path); err != nil {
		return err
	}

	return nil
}

// MkdirForFile helper creates the directory components of path.
func MkdirForFile(path string) error {
	return os.MkdirAll(filepath.Dir(path), DefaultDirectoryPermissions)
}
<commit_msg>internal/util/file: add function to get user & group id<commit_after>// Copyright 2015 CoreOS, Inc.
//
// 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 util

import (
	"encoding/hex"
	"hash"
	"io"
	"io/ioutil"
	"net/url"
	"os"
	"path/filepath"
	"strconv"

	"github.com/coreos/ignition/config/types"
	"github.com/coreos/ignition/internal/log"
	"github.com/coreos/ignition/internal/resource"
	internalUtil "github.com/coreos/ignition/internal/util"
)

const (
	DefaultDirectoryPermissions os.FileMode = 0755
	DefaultFilePermissions      os.FileMode = 0644
)

type FetchOp struct {
	Hash         hash.Hash
	Path         string
	Mode         os.FileMode
	Uid          int
	Gid          int
	Url          url.URL
	FetchOptions resource.FetchOptions
}

// newHashedReader returns a new ReadCloser that also writes to the provided hash.
func newHashedReader(reader io.ReadCloser, hasher hash.Hash) io.ReadCloser {
	return struct {
		io.Reader
		io.Closer
	}{
		Reader: io.TeeReader(reader, hasher),
		Closer: reader,
	}
}

// PrepareFetch converts a given logger, http client, and types.File into a
// FetchOp. This includes operations such as parsing the source URL, generating
// a hasher, and performing user/group name lookups. If an error is encountered,
// the issue will be logged and nil will be returned.
func (u Util) PrepareFetch(l *log.Logger, f types.File) *FetchOp {
	var err error
	var expectedSum []byte

	// explicitly ignoring the error here because the config should already be
	// validated by this point
	uri, _ := url.Parse(f.Contents.Source)

	hasher, err := GetHasher(f.Contents.Verification)
	if err != nil {
		l.Crit("Error verifying file %q: %v", f.Path, err)
		return nil
	}

	if hasher != nil {
		// explicitly ignoring the error here because the config should already
		// be validated by this point
		_, expectedSumString, _ := f.Contents.Verification.HashParts()
		expectedSum, err = hex.DecodeString(expectedSumString)
		if err != nil {
			l.Crit("Error parsing verification string %q: %v", expectedSumString, err)
			return nil
		}
	}

	f.User.ID, f.Group.ID = u.GetUserGroupID(l, f.User, f.Group)

	return &FetchOp{
		Path: f.Path,
		Hash: hasher,
		Mode: os.FileMode(f.Mode),
		Uid:  *f.User.ID,
		Gid:  *f.Group.ID,
		Url:  *uri,
		FetchOptions: resource.FetchOptions{
			Hash:        hasher,
			Compression: f.Contents.Compression,
			ExpectedSum: expectedSum,
		},
	}
}

func (u Util) WriteLink(s types.Link) error {
	path := u.JoinPath(s.Path)

	if err := MkdirForFile(path); err != nil {
		return err
	}

	if s.Hard {
		targetPath := u.JoinPath(s.Target)
		return os.Link(targetPath, path)
	}
	return os.Symlink(s.Target, path)
}

// PerformFetch performs a fetch operation generated by PrepareFetch, retrieving
// the file and writing it to disk. Any encountered errors are returned.
func (u Util) PerformFetch(f *FetchOp) error {
	var err error

	path := u.JoinPath(string(f.Path))

	if err := MkdirForFile(path); err != nil {
		return err
	}

	// Create a temporary file in the same directory to ensure it's on the same filesystem
	var tmp *os.File
	if tmp, err = ioutil.TempFile(filepath.Dir(path), "tmp"); err != nil {
		return err
	}

	defer func() {
		tmp.Close()
		if err != nil {
			os.Remove(tmp.Name())
		}
	}()

	err = u.Fetcher.Fetch(f.Url, tmp, f.FetchOptions)
	if err != nil {
		u.Crit("Error fetching file %q: %v", f.Path, err)
		return err
	}

	// XXX(vc): Note that we assume to be operating on the file we just wrote, this is only guaranteed
	// by using syscall.Fchown() and syscall.Fchmod()

	// Ensure the ownership and mode are as requested (since WriteFile can be affected by sticky bit)
	if err = os.Chown(tmp.Name(), f.Uid, f.Gid); err != nil {
		return err
	}

	if err = os.Chmod(tmp.Name(), f.Mode); err != nil {
		return err
	}

	if err = os.Rename(tmp.Name(), path); err != nil {
		return err
	}

	return nil
}

func (u Util) GetUserGroupID(l *log.Logger, user types.NodeUser, group types.NodeGroup) (*int, *int) {
	if user.Name != "" {
		usr, err := u.userLookup(user.Name)
		if err != nil {
			l.Crit("No such user %q: %v", user.Name, err)
			return nil, nil
		}
		uid, err := strconv.ParseInt(usr.Uid, 0, 0)
		if err != nil {
			l.Crit("Couldn't parse uid %q: %v", usr.Uid, err)
			return nil, nil
		}
		tmp := int(uid)
		user.ID = &tmp
	}
	if group.Name != "" {
		g, err := u.groupLookup(group.Name)
		if err != nil {
			l.Crit("No such group %q: %v", group.Name, err)
			return nil, nil
		}
		gid, err := strconv.ParseInt(g.Gid, 0, 0)
		if err != nil {
			l.Crit("Couldn't parse gid %q: %v", g.Gid, err)
			return nil, nil
		}
		tmp := int(gid)
		group.ID = &tmp
	}

	if user.ID == nil {
		user.ID = internalUtil.IntToPtr(0)
	}
	if group.ID == nil {
		group.ID = internalUtil.IntToPtr(0)
	}

	return user.ID, group.ID
}

// MkdirForFile helper creates the directory components of path.
func MkdirForFile(path string) error {
	return os.MkdirAll(filepath.Dir(path), DefaultDirectoryPermissions)
}
<|endoftext|>