<commit_before>package main

import (
    "bytes"
    "fmt"
    "github.com/coreos/go-etcd/etcd"
    "github.com/miekg/dns"
    "github.com/rcrowley/go-metrics"
    "net"
    "strconv"
    "strings"
    "sync"
    "time"
)

type Resolver struct {
    etcd        *etcd.Client
    etcdPrefix  string
}

// GetFromStorage looks up a key in etcd and returns a slice of nodes. It supports two storage structures;
//  - File:         /foo/bar/.A -> "value"
//  - Directory:    /foo/bar/.A/0 -> "value-0"
//                  /foo/bar/.A/1 -> "value-1"
func (r *Resolver) GetFromStorage(key string) (nodes []*etcd.Node, err error) {

    counter := metrics.GetOrRegisterCounter("resolver.etcd.query_count", metrics.DefaultRegistry)
    error_counter := metrics.GetOrRegisterCounter("resolver.etcd.query_error_count", metrics.DefaultRegistry)

    counter.Inc(1)
    debugMsg("Querying etcd for " + key)

    response, err := r.etcd.Get(r.etcdPrefix + key, false, true)
    if err != nil {
        error_counter.Inc(1)
        return
    }

    var findKeys func(node *etcd.Node)

    nodes = make([]*etcd.Node, 0)
    findKeys = func(node *etcd.Node) {
        if node.Dir == true {
            for _, subnode := range node.Nodes {
                findKeys(subnode)
            }
        } else {
            nodes = append(nodes, node)
        }
    }

    findKeys(response.Node)

    return
}

// Authority returns a dns.RR describing the know authority for the given
// domain. It will recurse up the domain structure to find an SOA record that
// matches.
func (r *Resolver) Authority(domain string) (soa *dns.SOA) {
    tree := strings.Split(domain, ".")
    for i, _ := range tree {
        subdomain := strings.Join(tree[i:], ".")

        // Check for an SOA entry
        answers, err := r.LookupAnswersForType(subdomain, dns.TypeSOA)
        if err != nil {
            return
        }

        if len(answers) == 1 {
            soa = answers[0].(*dns.SOA)
            soa.Serial = uint32(time.Now().Truncate(time.Hour).Unix())
            return
        }
    }

    // Maintain a counter for when we don't have an authority for a domain.
    missing_counter := metrics.GetOrRegisterCounter("resolver.authority.missing_soa", metrics.DefaultRegistry)
    missing_counter.Inc(1)

    return
}

// Lookup responds to DNS messages of type Query, with a dns message containing Answers.
// In the event that the query's value+type yields no known records, this falls back to
// querying the given nameservers instead.
func (r *Resolver) Lookup(req *dns.Msg) (msg *dns.Msg) {
    q := req.Question[0]

    msg = new(dns.Msg)
    msg.SetReply(req)
    msg.Authoritative = true
    msg.RecursionAvailable = false // We're a nameserver, no recursion for you!

    wait := sync.WaitGroup{}
    answers := make(chan dns.RR)
    errors := make(chan error)

    if q.Qclass == dns.ClassINET {
        r.AnswerQuestion(answers, errors, q, &wait)
    }

    // Spawn a goroutine to close the channel as soon as all of the things
    // are done. This allows us to ensure we'll wait for all workers to finish
    // but allows us to collect up answers concurrently.
    go func() {
        wait.Wait()

        // If we failed to find any answers, let's keep looking up the tree for
        // any wildcard domain entries.
        if len(answers) == 0 {
            parts := strings.Split(q.Name, ".")
            for level := 1; level < len(parts); level++ {
                domain := strings.Join(parts[level:], ".")
                if len(domain) > 1 {
                    question := dns.Question{
                        Name: "*." + dns.Fqdn(domain),
                        Qtype: q.Qtype,
                        Qclass: q.Qclass}

                    r.AnswerQuestion(answers, errors, question, &wait)

                    wait.Wait()
                    if len(answers) > 0 {
                        break;
                    }
                }
            }
        }

        debugMsg("Finished processing all goroutines, closing channels")
        close(answers)
        close(errors)
    }()

    miss_counter := metrics.GetOrRegisterCounter("resolver.answers.miss", metrics.DefaultRegistry)
    hit_counter := metrics.GetOrRegisterCounter("resolver.answers.hit", metrics.DefaultRegistry)
    error_counter := metrics.GetOrRegisterCounter("resolver.answers.error", metrics.DefaultRegistry)

    // Collect up all of the answers and any errors
    done := 0
    for done < 2 {
        select {
        case rr, ok := <-answers:
            if ok {
                rr.Header().Name = q.Name
                msg.Answer = append(msg.Answer, rr)
            } else {
                done++
            }
        case err, ok := <-errors:
            if ok {
                error_counter.Inc(1)
                // TODO(tarnfeld): Send special TXT records with a server error response code
                debugMsg("Error")
                debugMsg(err)
            } else {
                done++
            }
        }
    }

    // Send the correct authority records
    soa := r.Authority(q.Name)
    if len(msg.Answer) == 0 {
        miss_counter.Inc(1)
        msg.SetRcode(req, dns.RcodeNameError)
        if soa != nil {
            msg.Ns = []dns.RR{soa}
        } else {
            msg.Authoritative = false // No SOA? We're not authoritative
        }
    } else {
        hit_counter.Inc(1)
    }

    return
}

// AnswerQuestion takes two channels, one for answers and one for errors. It will answer the
// given question writing the answers as dns.RR structures, and any errors it encounters along
// the way. The function will return immediately, and spawn off a bunch of goroutines
// to do the work, when using this function one should use a WaitGroup to know when all work
// has been completed.
func (r *Resolver) AnswerQuestion(answers chan dns.RR, errors chan error, q dns.Question, wg *sync.WaitGroup) {

    typeStr := strings.ToLower(dns.TypeToString[q.Qtype])
    type_counter := metrics.GetOrRegisterCounter("resolver.answers.type." + typeStr, metrics.DefaultRegistry)
    type_counter.Inc(1)

    debugMsg("Answering question ", q)

    if q.Qtype == dns.TypeANY {
        wg.Add(len(converters))

        for rrType, _ := range converters {
            go func(rrType uint16) {
                defer func() { recover() }()
                defer wg.Done()

                results, err := r.LookupAnswersForType(q.Name, rrType)
                if err != nil {
                    errors <- err
                } else {
                    for _, answer := range results {
                        answers <- answer
                    }
                }
            }(rrType)
        }
    } else if _, ok := converters[q.Qtype]; ok {
        wg.Add(1)

        go func() {
            records, err := r.LookupAnswersForType(q.Name, q.Qtype)
            if err != nil {
                errors <- err
            } else {
                if len(records) > 0 {
                    for _, rr := range records {
                        answers <- rr
                    }
                } else {
                    cnames, err := r.LookupAnswersForType(q.Name, dns.TypeCNAME)
                    if err != nil {
                        errors <- err
                    } else {
                        if len(cnames) > 1 {
                            errors <- &RecordValueError{
                                Message: "Multiple CNAME records is invalid",
                                AttemptedType: dns.TypeCNAME}
                        } else if len(cnames) > 0 {
                            answers <- cnames[0]
                        }
                    }
                }
            }

            wg.Done()
        }()
    }
}

func (r *Resolver) LookupAnswersForType(name string, rrType uint16) (answers []dns.RR, err error) {
    name = strings.ToLower(name)

    typeStr := dns.TypeToString[rrType]
    nodes, err := r.GetFromStorage(nameToKey(name, "/." + typeStr))

    if err != nil {
        if e, ok := err.(*etcd.EtcdError); ok {
            if e.ErrorCode == 100 {
                return answers, nil
            }
        }

        return
    }

    answers = make([]dns.RR, len(nodes))
    for i, node := range nodes {

        // TODO(tarnfeld): TTL 0 - make this configurable
        header := dns.RR_Header{Name: name, Class: dns.ClassINET, Rrtype: rrType, Ttl: 0}
        answer, err := converters[rrType](node, header)

        if err != nil {
            debugMsg("Error converting type: ", err)
            return nil, err
        }

        answers[i] = answer
    }

    return
}

// nameToKey returns a string representing the etcd version of a domain, replacing dots with slashes
// and reversing it (foo.net. -> /net/foo)
func nameToKey(name string, suffix string) string {
    segments := strings.Split(name, ".")

    var keyBuffer bytes.Buffer
    for i := len(segments) - 1; i >= 0; i-- {
        if len(segments[i]) > 0 {
            keyBuffer.WriteString("/")
            keyBuffer.WriteString(segments[i])
        }
    }

    keyBuffer.WriteString(suffix)
    return keyBuffer.String()
}

// Map of conversion functions that turn individual etcd nodes into dns.RR answers
var converters = map[uint16]func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {

    dns.TypeA: func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {

        ip := net.ParseIP(node.Value)
        if ip == nil {
            err = &NodeConversionError{
                Node: node,
                Message: fmt.Sprintf("Failed to parse %s as IP Address", node.Value),
                AttemptedType: dns.TypeA,
            }
        } else if ip.To4() == nil {
            err = &NodeConversionError{
                Node: node,
                Message: fmt.Sprintf("Value %s isn't an IPv4 address", node.Value),
                AttemptedType: dns.TypeA,
            }
        } else {
            rr = &dns.A{header, ip}
        }

        return
    },

    dns.TypeAAAA: func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {

        ip := net.ParseIP(node.Value)
        if ip == nil {
            err = &NodeConversionError{
                Node: node,
                Message: fmt.Sprintf("Failed to parse IP Address %s", node.Value),
                AttemptedType: dns.TypeAAAA}
        } else if ip.To16() == nil {
            err = &NodeConversionError{
                Node: node,
                Message: fmt.Sprintf("Value %s isn't an IPv6 address", node.Value),
                AttemptedType: dns.TypeA}
        } else {
            rr = &dns.AAAA{header, ip}
        }
        return
    },

    dns.TypeTXT: func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {
        rr = &dns.TXT{header, []string{node.Value}}
        return
    },

    dns.TypeCNAME: func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {
        rr = &dns.CNAME{header, dns.Fqdn(node.Value)}
        return
    },

    dns.TypeNS: func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {
        rr = &dns.NS{header, dns.Fqdn(node.Value)}
        return
    },

    dns.TypePTR: func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {
        labels, ok := dns.IsDomainName(node.Value)

        if (ok && labels > 0) {
            rr = &dns.PTR{header, dns.Fqdn(node.Value)}
        } else {
            err = &NodeConversionError{
                Node: node,
                Message: fmt.Sprintf("Value '%s' isn't a valid domain name", node.Value),
                AttemptedType: dns.TypePTR}
        }
        return
    },

    dns.TypeSOA: func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {
        parts := strings.SplitN(node.Value, "\\t", 6)

        if len(parts) < 6 {
            err = &NodeConversionError{
                Node: node,
                Message: fmt.Sprintf("Value %s isn't valid for SOA", node.Value),
                AttemptedType: dns.TypeSOA}
        } else {
            refresh, err := strconv.ParseUint(parts[2], 10, 32)
            if err != nil {
                return nil, err
            }

            retry, err := strconv.ParseUint(parts[3], 10, 32)
            if err != nil {
                return nil, err
            }

            expire, err := strconv.ParseUint(parts[4], 10, 32)
            if err != nil {
                return nil, err
            }

            minttl, err := strconv.ParseUint(parts[5], 10, 32)
            if err != nil {
                return nil, err
            }

            rr = &dns.SOA{
                Hdr:     header,
                Ns:      dns.Fqdn(parts[0]),
                Mbox:    dns.Fqdn(parts[1]),
                Refresh: uint32(refresh),
                Retry:   uint32(retry),
                Expire:  uint32(expire),
                Minttl:  uint32(minttl)}
        }

        return
    },
}
<commit_msg>Add support for SRV records<commit_after>package main

import (
    "bytes"
    "fmt"
    "github.com/coreos/go-etcd/etcd"
    "github.com/miekg/dns"
    "github.com/rcrowley/go-metrics"
    "net"
    "strconv"
    "strings"
    "sync"
    "time"
)

type Resolver struct {
    etcd        *etcd.Client
    etcdPrefix  string
}

// GetFromStorage looks up a key in etcd and returns a slice of nodes. It supports two storage structures;
//  - File:         /foo/bar/.A -> "value"
//  - Directory:    /foo/bar/.A/0 -> "value-0"
//                  /foo/bar/.A/1 -> "value-1"
func (r *Resolver) GetFromStorage(key string) (nodes []*etcd.Node, err error) {

    counter := metrics.GetOrRegisterCounter("resolver.etcd.query_count", metrics.DefaultRegistry)
    error_counter := metrics.GetOrRegisterCounter("resolver.etcd.query_error_count", metrics.DefaultRegistry)

    counter.Inc(1)
    debugMsg("Querying etcd for " + key)

    response, err := r.etcd.Get(r.etcdPrefix + key, false, true)
    if err != nil {
        error_counter.Inc(1)
        return
    }

    var findKeys func(node *etcd.Node)

    nodes = make([]*etcd.Node, 0)
    findKeys = func(node *etcd.Node) {
        if node.Dir == true {
            for _, subnode := range node.Nodes {
                findKeys(subnode)
            }
        } else {
            nodes = append(nodes, node)
        }
    }

    findKeys(response.Node)

    return
}

// Authority returns a dns.RR describing the know authority for the given
// domain. It will recurse up the domain structure to find an SOA record that
// matches.
func (r *Resolver) Authority(domain string) (soa *dns.SOA) {
    tree := strings.Split(domain, ".")
    for i, _ := range tree {
        subdomain := strings.Join(tree[i:], ".")

        // Check for an SOA entry
        answers, err := r.LookupAnswersForType(subdomain, dns.TypeSOA)
        if err != nil {
            return
        }

        if len(answers) == 1 {
            soa = answers[0].(*dns.SOA)
            soa.Serial = uint32(time.Now().Truncate(time.Hour).Unix())
            return
        }
    }

    // Maintain a counter for when we don't have an authority for a domain.
    missing_counter := metrics.GetOrRegisterCounter("resolver.authority.missing_soa", metrics.DefaultRegistry)
    missing_counter.Inc(1)

    return
}

// Lookup responds to DNS messages of type Query, with a dns message containing Answers.
// In the event that the query's value+type yields no known records, this falls back to
// querying the given nameservers instead.
func (r *Resolver) Lookup(req *dns.Msg) (msg *dns.Msg) {
    q := req.Question[0]

    msg = new(dns.Msg)
    msg.SetReply(req)
    msg.Authoritative = true
    msg.RecursionAvailable = false // We're a nameserver, no recursion for you!

    wait := sync.WaitGroup{}
    answers := make(chan dns.RR)
    errors := make(chan error)

    if q.Qclass == dns.ClassINET {
        r.AnswerQuestion(answers, errors, q, &wait)
    }

    // Spawn a goroutine to close the channel as soon as all of the things
    // are done. This allows us to ensure we'll wait for all workers to finish
    // but allows us to collect up answers concurrently.
    go func() {
        wait.Wait()

        // If we failed to find any answers, let's keep looking up the tree for
        // any wildcard domain entries.
        if len(answers) == 0 {
            parts := strings.Split(q.Name, ".")
            for level := 1; level < len(parts); level++ {
                domain := strings.Join(parts[level:], ".")
                if len(domain) > 1 {
                    question := dns.Question{
                        Name: "*." + dns.Fqdn(domain),
                        Qtype: q.Qtype,
                        Qclass: q.Qclass}

                    r.AnswerQuestion(answers, errors, question, &wait)

                    wait.Wait()
                    if len(answers) > 0 {
                        break;
                    }
                }
            }
        }

        debugMsg("Finished processing all goroutines, closing channels")
        close(answers)
        close(errors)
    }()

    miss_counter := metrics.GetOrRegisterCounter("resolver.answers.miss", metrics.DefaultRegistry)
    hit_counter := metrics.GetOrRegisterCounter("resolver.answers.hit", metrics.DefaultRegistry)
    error_counter := metrics.GetOrRegisterCounter("resolver.answers.error", metrics.DefaultRegistry)

    // Collect up all of the answers and any errors
    done := 0
    for done < 2 {
        select {
        case rr, ok := <-answers:
            if ok {
                rr.Header().Name = q.Name
                msg.Answer = append(msg.Answer, rr)
            } else {
                done++
            }
        case err, ok := <-errors:
            if ok {
                error_counter.Inc(1)
                // TODO(tarnfeld): Send special TXT records with a server error response code
                debugMsg("Error")
                debugMsg(err)
            } else {
                done++
            }
        }
    }

    // Send the correct authority records
    soa := r.Authority(q.Name)
    if len(msg.Answer) == 0 {
        miss_counter.Inc(1)
        msg.SetRcode(req, dns.RcodeNameError)
        if soa != nil {
            msg.Ns = []dns.RR{soa}
        } else {
            msg.Authoritative = false // No SOA? We're not authoritative
        }
    } else {
        hit_counter.Inc(1)
    }

    return
}

// AnswerQuestion takes two channels, one for answers and one for errors. It will answer the
// given question writing the answers as dns.RR structures, and any errors it encounters along
// the way. The function will return immediately, and spawn off a bunch of goroutines
// to do the work, when using this function one should use a WaitGroup to know when all work
// has been completed.
func (r *Resolver) AnswerQuestion(answers chan dns.RR, errors chan error, q dns.Question, wg *sync.WaitGroup) {

    typeStr := strings.ToLower(dns.TypeToString[q.Qtype])
    type_counter := metrics.GetOrRegisterCounter("resolver.answers.type." + typeStr, metrics.DefaultRegistry)
    type_counter.Inc(1)

    debugMsg("Answering question ", q)

    if q.Qtype == dns.TypeANY {
        wg.Add(len(converters))

        for rrType, _ := range converters {
            go func(rrType uint16) {
                defer func() { recover() }()
                defer wg.Done()

                results, err := r.LookupAnswersForType(q.Name, rrType)
                if err != nil {
                    errors <- err
                } else {
                    for _, answer := range results {
                        answers <- answer
                    }
                }
            }(rrType)
        }
    } else if _, ok := converters[q.Qtype]; ok {
        wg.Add(1)

        go func() {
            records, err := r.LookupAnswersForType(q.Name, q.Qtype)
            if err != nil {
                errors <- err
            } else {
                if len(records) > 0 {
                    for _, rr := range records {
                        answers <- rr
                    }
                } else {
                    cnames, err := r.LookupAnswersForType(q.Name, dns.TypeCNAME)
                    if err != nil {
                        errors <- err
                    } else {
                        if len(cnames) > 1 {
                            errors <- &RecordValueError{
                                Message: "Multiple CNAME records is invalid",
                                AttemptedType: dns.TypeCNAME}
                        } else if len(cnames) > 0 {
                            answers <- cnames[0]
                        }
                    }
                }
            }

            wg.Done()
        }()
    }
}

func (r *Resolver) LookupAnswersForType(name string, rrType uint16) (answers []dns.RR, err error) {
    name = strings.ToLower(name)

    typeStr := dns.TypeToString[rrType]
    nodes, err := r.GetFromStorage(nameToKey(name, "/." + typeStr))

    if err != nil {
        if e, ok := err.(*etcd.EtcdError); ok {
            if e.ErrorCode == 100 {
                return answers, nil
            }
        }

        return
    }

    answers = make([]dns.RR, len(nodes))
    for i, node := range nodes {

        // TODO(tarnfeld): TTL 0 - make this configurable
        header := dns.RR_Header{Name: name, Class: dns.ClassINET, Rrtype: rrType, Ttl: 0}
        answer, err := converters[rrType](node, header)

        if err != nil {
            debugMsg("Error converting type: ", err)
            return nil, err
        }

        answers[i] = answer
    }

    return
}

// nameToKey returns a string representing the etcd version of a domain, replacing dots with slashes
// and reversing it (foo.net. -> /net/foo)
func nameToKey(name string, suffix string) string {
    segments := strings.Split(name, ".")

    var keyBuffer bytes.Buffer
    for i := len(segments) - 1; i >= 0; i-- {
        if len(segments[i]) > 0 {
            keyBuffer.WriteString("/")
            keyBuffer.WriteString(segments[i])
        }
    }

    keyBuffer.WriteString(suffix)
    return keyBuffer.String()
}

// Map of conversion functions that turn individual etcd nodes into dns.RR answers
var converters = map[uint16]func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {

    dns.TypeA: func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {

        ip := net.ParseIP(node.Value)
        if ip == nil {
            err = &NodeConversionError{
                Node: node,
                Message: fmt.Sprintf("Failed to parse %s as IP Address", node.Value),
                AttemptedType: dns.TypeA,
            }
        } else if ip.To4() == nil {
            err = &NodeConversionError{
                Node: node,
                Message: fmt.Sprintf("Value %s isn't an IPv4 address", node.Value),
                AttemptedType: dns.TypeA,
            }
        } else {
            rr = &dns.A{header, ip}
        }

        return
    },

    dns.TypeAAAA: func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {

        ip := net.ParseIP(node.Value)
        if ip == nil {
            err = &NodeConversionError{
                Node: node,
                Message: fmt.Sprintf("Failed to parse IP Address %s", node.Value),
                AttemptedType: dns.TypeAAAA}
        } else if ip.To16() == nil {
            err = &NodeConversionError{
                Node: node,
                Message: fmt.Sprintf("Value %s isn't an IPv6 address", node.Value),
                AttemptedType: dns.TypeA}
        } else {
            rr = &dns.AAAA{header, ip}
        }
        return
    },

    dns.TypeTXT: func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {
        rr = &dns.TXT{header, []string{node.Value}}
        return
    },

    dns.TypeCNAME: func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {
        rr = &dns.CNAME{header, dns.Fqdn(node.Value)}
        return
    },

    dns.TypeNS: func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {
        rr = &dns.NS{header, dns.Fqdn(node.Value)}
        return
    },

    dns.TypePTR: func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {
        labels, ok := dns.IsDomainName(node.Value)

        if (ok && labels > 0) {
            rr = &dns.PTR{header, dns.Fqdn(node.Value)}
        } else {
            err = &NodeConversionError{
                Node: node,
                Message: fmt.Sprintf("Value '%s' isn't a valid domain name", node.Value),
                AttemptedType: dns.TypePTR}
        }
        return
    },

    dns.TypeSRV: func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {
        parts := strings.SplitN(node.Value, "\\t", 4)

        if len(parts) != 4 {
            err = &NodeConversionError{
                Node: node,
                Message: fmt.Sprintf("Value %s isn't valid for SRV", node.Value),
                AttemptedType: dns.TypeSRV}
        } else {

            priority, err := strconv.ParseUint(parts[0], 10, 16)
            if err != nil {
                return nil, err
            }

            weight, err := strconv.ParseUint(parts[1], 10, 16)
            if err != nil {
                return nil, err
            }

            port, err := strconv.ParseUint(parts[2], 10, 16)
            if err != nil {
                return nil, err
            }

            labels, ok := dns.IsDomainName(parts[3])

            if (!ok || labels == 0) {
                err = &NodeConversionError{
                    Node: node,
                    Message: fmt.Sprintf("Value '%s' isn't a valid domain name", parts[3]),
                    AttemptedType: dns.TypeSRV}
                return nil, err
            }

            target := dns.Fqdn(parts[3])

            rr = &dns.SRV{
                header,
                uint16(priority),
                uint16(weight),
                uint16(port),
                target}
        }
        return
    },

    dns.TypeSOA: func (node *etcd.Node, header dns.RR_Header) (rr dns.RR, err error) {
        parts := strings.SplitN(node.Value, "\\t", 6)

        if len(parts) < 6 {
            err = &NodeConversionError{
                Node: node,
                Message: fmt.Sprintf("Value %s isn't valid for SOA", node.Value),
                AttemptedType: dns.TypeSOA}
        } else {
            refresh, err := strconv.ParseUint(parts[2], 10, 32)
            if err != nil {
                return nil, err
            }

            retry, err := strconv.ParseUint(parts[3], 10, 32)
            if err != nil {
                return nil, err
            }

            expire, err := strconv.ParseUint(parts[4], 10, 32)
            if err != nil {
                return nil, err
            }

            minttl, err := strconv.ParseUint(parts[5], 10, 32)
            if err != nil {
                return nil, err
            }

            rr = &dns.SOA{
                Hdr:     header,
                Ns:      dns.Fqdn(parts[0]),
                Mbox:    dns.Fqdn(parts[1]),
                Refresh: uint32(refresh),
                Retry:   uint32(retry),
                Expire:  uint32(expire),
                Minttl:  uint32(minttl)}
        }

        return
    },
}
<|endoftext|>


<commit_before>package restful

// Copyright 2013 Ernest Micklei. All rights reserved.
// Use of this source code is governed by a license
// that can be found in the LICENSE file.

import (
	"encoding/json"
	"encoding/xml"
	"net/http"
	"strings"
)

// DEPRECATED, use DefaultResponseContentType(mime)
var DefaultResponseMimeType string

//PrettyPrintResponses controls the indentation feature of XML and JSON
//serialization in the response methods WriteEntity, WriteAsJson, and
//WriteAsXml.
var PrettyPrintResponses = true

// Response is a wrapper on the actual http ResponseWriter
// It provides several convenience methods to prepare and write response content.
type Response struct {
	http.ResponseWriter
	requestAccept string   // mime-type what the Http Request says it wants to receive
	routeProduces []string // mime-types what the Route says it can produce
	statusCode    int      // HTTP status code that has been written explicity (if zero then net/http has written 200)
	contentLength int      // number of bytes written for the response body
	prettyPrint   bool     // controls the indentation feature of XML and JSON serialization. It is initialized using var PrettyPrintResponses.
}

// Creates a new response based on a http ResponseWriter.
func NewResponse(httpWriter http.ResponseWriter) *Response {
	return &Response{httpWriter, "", []string{}, http.StatusOK, 0, PrettyPrintResponses} // empty content-types
}

// If Accept header matching fails, fall back to this type, otherwise
// a "406: Not Acceptable" response is returned.
// Valid values are restful.MIME_JSON and restful.MIME_XML
// Example:
// 	restful.DefaultResponseContentType(restful.MIME_JSON)
func DefaultResponseContentType(mime string) {
	DefaultResponseMimeType = mime
}

// InternalServerError writes the StatusInternalServerError header.
// DEPRECATED, use WriteErrorString(http.StatusInternalServerError,reason)
func (r Response) InternalServerError() Response {
	r.WriteHeader(http.StatusInternalServerError)
	return r
}

// PrettyPrint changes whether this response must produce pretty (line-by-line, indented) JSON or XML output.
func (r *Response) PrettyPrint(bePretty bool) {
	r.prettyPrint = bePretty
}

// AddHeader is a shortcut for .Header().Add(header,value)
func (r Response) AddHeader(header string, value string) Response {
	r.Header().Add(header, value)
	return r
}

// SetRequestAccepts tells the response what Mime-type(s) the HTTP request said it wants to accept. Exposed for testing.
func (r *Response) SetRequestAccepts(mime string) {
	r.requestAccept = mime
}

// WriteEntity marshals the value using the representation denoted by the Accept Header (XML or JSON)
// If no Accept header is specified (or */*) then return the Content-Type as specified by the first in the Route.Produces.
// If an Accept header is specified then return the Content-Type as specified by the first in the Route.Produces that is matched with the Accept header.
// If the value is nil then nothing is written. You may want to call WriteHeader(http.StatusNotFound) instead.
// Current implementation ignores any q-parameters in the Accept Header.
func (r *Response) WriteEntity(value interface{}) error {
	if value == nil { // do not write a nil representation
		return nil
	}
	for _, qualifiedMime := range strings.Split(r.requestAccept, ",") {
		mime := strings.Trim(strings.Split(qualifiedMime, ";")[0], " ")
		if 0 == len(mime) || mime == "*/*" {
			for _, each := range r.routeProduces {
				if MIME_JSON == each {
					return r.WriteAsJson(value)
				}
				if MIME_XML == each {
					return r.WriteAsXml(value)
				}
			}
		} else { // mime is not blank; see if we have a match in Produces
			for _, each := range r.routeProduces {
				if mime == each {
					if MIME_JSON == each {
						return r.WriteAsJson(value)
					}
					if MIME_XML == each {
						return r.WriteAsXml(value)
					}
				}
			}
		}
	}
	if DefaultResponseMimeType == MIME_JSON {
		return r.WriteAsJson(value)
	} else if DefaultResponseMimeType == MIME_XML {
		return r.WriteAsXml(value)
	} else {
		if trace {
			traceLogger.Printf("mismatch in mime-types and no defaults; (http)Accept=%v,(route)Produces=%v\n", r.requestAccept, r.routeProduces)
		}
		r.WriteHeader(http.StatusNotAcceptable) // for recording only
		r.ResponseWriter.WriteHeader(http.StatusNotAcceptable)
		if _, err := r.Write([]byte("406: Not Acceptable")); err != nil {
			return err
		}
	}
	return nil
}

// WriteAsXml is a convenience method for writing a value in xml (requires Xml tags on the value)
func (r *Response) WriteAsXml(value interface{}) error {
	var output []byte
	var err error

	if value == nil { // do not write a nil representation
		return nil
	}
	if r.prettyPrint {
		output, err = xml.MarshalIndent(value, " ", " ")
	} else {
		output, err = xml.Marshal(value)
	}

	if err != nil {
		return r.WriteError(http.StatusInternalServerError, err)
	}
	r.Header().Set(HEADER_ContentType, MIME_XML)
	if r.statusCode > 0 { // a WriteHeader was intercepted
		r.ResponseWriter.WriteHeader(r.statusCode)
	}
	_, err = r.Write([]byte(xml.Header))
	if err != nil {
		return err
	}
	if _, err = r.Write(output); err != nil {
		return err
	}
	return nil
}

// WriteAsJson is a convenience method for writing a value in json
func (r *Response) WriteAsJson(value interface{}) error {
	return r.WriteJson(value, MIME_JSON) // no charset
}

// WriteJson is a convenience method for writing a value in Json with a given Content-Type
func (r *Response) WriteJson(value interface{}, contentType string) error {
	var output []byte
	var err error

	if value == nil { // do not write a nil representation
		return nil
	}
	if r.prettyPrint {
		output, err = json.MarshalIndent(value, " ", " ")
	} else {
		output, err = json.Marshal(value)
	}

	if err != nil {
		return r.WriteErrorString(http.StatusInternalServerError, err.Error())
	}
	r.Header().Set(HEADER_ContentType, contentType)
	if r.statusCode > 0 { // a WriteHeader was intercepted
		r.ResponseWriter.WriteHeader(r.statusCode)
	}
	if _, err = r.Write(output); err != nil {
		return err
	}
	return nil
}

// WriteError write the http status and the error string on the response.
func (r *Response) WriteError(httpStatus int, err error) error {
	return r.WriteErrorString(httpStatus, err.Error())
}

// WriteServiceError is a convenience method for a responding with a ServiceError and a status
func (r *Response) WriteServiceError(httpStatus int, err ServiceError) error {
	r.WriteHeader(httpStatus) // for recording only
	return r.WriteEntity(err)
}

// WriteErrorString is a convenience method for an error status with the actual error
func (r *Response) WriteErrorString(status int, errorReason string) error {
	r.statusCode = status // for recording only
	r.ResponseWriter.WriteHeader(status)
	if _, err := r.Write([]byte(errorReason)); err != nil {
		return err
	}
	return nil
}

// WriteHeader is overridden to remember the Status Code that has been written.
// Note that using this method, the status value is only written when
// - calling WriteEntity,
// - or directly calling WriteAsXml or WriteAsJson,
// - or if the status is one for which no response is allowed (i.e.,
//   204 (http.StatusNoContent) or 304 (http.StatusNotModified))
func (r *Response) WriteHeader(httpStatus int) {
	r.statusCode = httpStatus
	// if 201,204,304 then WriteEntity will not be called so we need to pass this code
	if http.StatusNoContent == httpStatus ||
		http.StatusNotModified == httpStatus ||
		http.StatusPartialContent == httpStatus {
		r.ResponseWriter.WriteHeader(httpStatus)
	}
}

// StatusCode returns the code that has been written using WriteHeader.
func (r Response) StatusCode() int {
	if 0 == r.statusCode {
		// no status code has been written yet; assume OK
		return http.StatusOK
	}
	return r.statusCode
}

// Write writes the data to the connection as part of an HTTP reply.
// Write is part of http.ResponseWriter interface.
func (r *Response) Write(bytes []byte) (int, error) {
	written, err := r.ResponseWriter.Write(bytes)
	r.contentLength += written
	return written, err
}

// ContentLength returns the number of bytes written for the response content.
// Note that this value is only correct if all data is written through the Response using its Write* methods.
// Data written directly using the underlying http.ResponseWriter is not accounted for.
func (r Response) ContentLength() int {
	return r.contentLength
}

// CloseNotify is part of http.CloseNotifier interface
func (r Response) CloseNotify() <-chan bool {
	return r.ResponseWriter.(http.CloseNotifier).CloseNotify()
}
<commit_msg>update WriteHeader comment<commit_after>package restful

// Copyright 2013 Ernest Micklei. All rights reserved.
// Use of this source code is governed by a license
// that can be found in the LICENSE file.

import (
	"encoding/json"
	"encoding/xml"
	"net/http"
	"strings"
)

// DEPRECATED, use DefaultResponseContentType(mime)
var DefaultResponseMimeType string

//PrettyPrintResponses controls the indentation feature of XML and JSON
//serialization in the response methods WriteEntity, WriteAsJson, and
//WriteAsXml.
var PrettyPrintResponses = true

// Response is a wrapper on the actual http ResponseWriter
// It provides several convenience methods to prepare and write response content.
type Response struct {
	http.ResponseWriter
	requestAccept string   // mime-type what the Http Request says it wants to receive
	routeProduces []string // mime-types what the Route says it can produce
	statusCode    int      // HTTP status code that has been written explicity (if zero then net/http has written 200)
	contentLength int      // number of bytes written for the response body
	prettyPrint   bool     // controls the indentation feature of XML and JSON serialization. It is initialized using var PrettyPrintResponses.
}

// Creates a new response based on a http ResponseWriter.
func NewResponse(httpWriter http.ResponseWriter) *Response {
	return &Response{httpWriter, "", []string{}, http.StatusOK, 0, PrettyPrintResponses} // empty content-types
}

// If Accept header matching fails, fall back to this type, otherwise
// a "406: Not Acceptable" response is returned.
// Valid values are restful.MIME_JSON and restful.MIME_XML
// Example:
// 	restful.DefaultResponseContentType(restful.MIME_JSON)
func DefaultResponseContentType(mime string) {
	DefaultResponseMimeType = mime
}

// InternalServerError writes the StatusInternalServerError header.
// DEPRECATED, use WriteErrorString(http.StatusInternalServerError,reason)
func (r Response) InternalServerError() Response {
	r.WriteHeader(http.StatusInternalServerError)
	return r
}

// PrettyPrint changes whether this response must produce pretty (line-by-line, indented) JSON or XML output.
func (r *Response) PrettyPrint(bePretty bool) {
	r.prettyPrint = bePretty
}

// AddHeader is a shortcut for .Header().Add(header,value)
func (r Response) AddHeader(header string, value string) Response {
	r.Header().Add(header, value)
	return r
}

// SetRequestAccepts tells the response what Mime-type(s) the HTTP request said it wants to accept. Exposed for testing.
func (r *Response) SetRequestAccepts(mime string) {
	r.requestAccept = mime
}

// WriteEntity marshals the value using the representation denoted by the Accept Header (XML or JSON)
// If no Accept header is specified (or */*) then return the Content-Type as specified by the first in the Route.Produces.
// If an Accept header is specified then return the Content-Type as specified by the first in the Route.Produces that is matched with the Accept header.
// If the value is nil then nothing is written. You may want to call WriteHeader(http.StatusNotFound) instead.
// Current implementation ignores any q-parameters in the Accept Header.
func (r *Response) WriteEntity(value interface{}) error {
	if value == nil { // do not write a nil representation
		return nil
	}
	for _, qualifiedMime := range strings.Split(r.requestAccept, ",") {
		mime := strings.Trim(strings.Split(qualifiedMime, ";")[0], " ")
		if 0 == len(mime) || mime == "*/*" {
			for _, each := range r.routeProduces {
				if MIME_JSON == each {
					return r.WriteAsJson(value)
				}
				if MIME_XML == each {
					return r.WriteAsXml(value)
				}
			}
		} else { // mime is not blank; see if we have a match in Produces
			for _, each := range r.routeProduces {
				if mime == each {
					if MIME_JSON == each {
						return r.WriteAsJson(value)
					}
					if MIME_XML == each {
						return r.WriteAsXml(value)
					}
				}
			}
		}
	}
	if DefaultResponseMimeType == MIME_JSON {
		return r.WriteAsJson(value)
	} else if DefaultResponseMimeType == MIME_XML {
		return r.WriteAsXml(value)
	} else {
		if trace {
			traceLogger.Printf("mismatch in mime-types and no defaults; (http)Accept=%v,(route)Produces=%v\n", r.requestAccept, r.routeProduces)
		}
		r.WriteHeader(http.StatusNotAcceptable) // for recording only
		r.ResponseWriter.WriteHeader(http.StatusNotAcceptable)
		if _, err := r.Write([]byte("406: Not Acceptable")); err != nil {
			return err
		}
	}
	return nil
}

// WriteAsXml is a convenience method for writing a value in xml (requires Xml tags on the value)
func (r *Response) WriteAsXml(value interface{}) error {
	var output []byte
	var err error

	if value == nil { // do not write a nil representation
		return nil
	}
	if r.prettyPrint {
		output, err = xml.MarshalIndent(value, " ", " ")
	} else {
		output, err = xml.Marshal(value)
	}

	if err != nil {
		return r.WriteError(http.StatusInternalServerError, err)
	}
	r.Header().Set(HEADER_ContentType, MIME_XML)
	if r.statusCode > 0 { // a WriteHeader was intercepted
		r.ResponseWriter.WriteHeader(r.statusCode)
	}
	_, err = r.Write([]byte(xml.Header))
	if err != nil {
		return err
	}
	if _, err = r.Write(output); err != nil {
		return err
	}
	return nil
}

// WriteAsJson is a convenience method for writing a value in json
func (r *Response) WriteAsJson(value interface{}) error {
	return r.WriteJson(value, MIME_JSON) // no charset
}

// WriteJson is a convenience method for writing a value in Json with a given Content-Type
func (r *Response) WriteJson(value interface{}, contentType string) error {
	var output []byte
	var err error

	if value == nil { // do not write a nil representation
		return nil
	}
	if r.prettyPrint {
		output, err = json.MarshalIndent(value, " ", " ")
	} else {
		output, err = json.Marshal(value)
	}

	if err != nil {
		return r.WriteErrorString(http.StatusInternalServerError, err.Error())
	}
	r.Header().Set(HEADER_ContentType, contentType)
	if r.statusCode > 0 { // a WriteHeader was intercepted
		r.ResponseWriter.WriteHeader(r.statusCode)
	}
	if _, err = r.Write(output); err != nil {
		return err
	}
	return nil
}

// WriteError write the http status and the error string on the response.
func (r *Response) WriteError(httpStatus int, err error) error {
	return r.WriteErrorString(httpStatus, err.Error())
}

// WriteServiceError is a convenience method for a responding with a ServiceError and a status
func (r *Response) WriteServiceError(httpStatus int, err ServiceError) error {
	r.WriteHeader(httpStatus) // for recording only
	return r.WriteEntity(err)
}

// WriteErrorString is a convenience method for an error status with the actual error
func (r *Response) WriteErrorString(status int, errorReason string) error {
	r.statusCode = status // for recording only
	r.ResponseWriter.WriteHeader(status)
	if _, err := r.Write([]byte(errorReason)); err != nil {
		return err
	}
	return nil
}

// WriteHeader is overridden to remember the Status Code that has been written.
// Note that using this method, the status value is only written when
// - calling WriteEntity,
// - or directly calling WriteAsXml or WriteAsJson,
// - or if the status is one for which no response is allowed (i.e.,
//   204 (http.StatusNoContent) or 304 (http.StatusNotModified) or 206 (http.StatusPartialContent)
func (r *Response) WriteHeader(httpStatus int) {
	r.statusCode = httpStatus
	// if 206,204,304 then WriteEntity will not be called so we need to pass this code
	if http.StatusNoContent == httpStatus ||
		http.StatusNotModified == httpStatus ||
		http.StatusPartialContent == httpStatus {
		r.ResponseWriter.WriteHeader(httpStatus)
	}
}

// StatusCode returns the code that has been written using WriteHeader.
func (r Response) StatusCode() int {
	if 0 == r.statusCode {
		// no status code has been written yet; assume OK
		return http.StatusOK
	}
	return r.statusCode
}

// Write writes the data to the connection as part of an HTTP reply.
// Write is part of http.ResponseWriter interface.
func (r *Response) Write(bytes []byte) (int, error) {
	written, err := r.ResponseWriter.Write(bytes)
	r.contentLength += written
	return written, err
}

// ContentLength returns the number of bytes written for the response content.
// Note that this value is only correct if all data is written through the Response using its Write* methods.
// Data written directly using the underlying http.ResponseWriter is not accounted for.
func (r Response) ContentLength() int {
	return r.contentLength
}

// CloseNotify is part of http.CloseNotifier interface
func (r Response) CloseNotify() <-chan bool {
	return r.ResponseWriter.(http.CloseNotifier).CloseNotify()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"

	"github.com/justincampbell/forecast/v2"
)

var conditionIcons = map[string]string{
	"clear-day":           "☀",
	"clear-night":         "🌙",
	"cloudy":              "☁",
	"fog":                 "🌁",
	"partly-cloudy-day":   "⛅",
	"partly-cloudy-night": "🌙",
	"rain":                "☔",
	"sleet":               "❄ ☔",
	"snow":                "❄",
	"wind":                "🍃",
}

var maxCacheAge, _ = time.ParseDuration("1h")

func main() {
	coordinates := flag.String("coordinates", "39.95,-75.1667", "the coordinates, expressed as latitude,longitude")
	tmpDir := flag.String("tmpdir", os.TempDir(), "the directory to use to store cached responses")
	key := flag.String("key", os.Getenv("FORECAST_IO_API_KEY"), "your forecast.io API key")

	if *key == "" {
		exitWith("Please provide your forecast.io API key with -key, or set FORECAST_IO_API_KEY", 1)
	}

	flag.Parse()

	coordinateParts := strings.Split(*coordinates, ",")
	var latitude string
	var longitude string

	if len(coordinateParts) != 2 {
		exitWith("You must specify latitude and longitude like so: 39.95,-75.1667", 1)
	} else {
		latitude = coordinateParts[0]
		longitude = coordinateParts[1]
	}

	var cacheFilename = fmt.Sprintf("emoji-weather-%s-%s.json", latitude, longitude)
	var cacheFile = path.Join(*tmpDir, cacheFilename)

	var json []byte
	var err error

	if isCacheStale(cacheFile) {
		json, err = getForecast(*key, latitude, longitude)
		check(err)

		err = writeCache(cacheFile, json)
		check(err)
	} else {
		json, err = ioutil.ReadFile(cacheFile)
		check(err)
	}

	fmt.Println(formatConditions(extractConditionFromJSON(json)))
}

func isCacheStale(cacheFile string) bool {
	stat, err := os.Stat(cacheFile)

	return os.IsNotExist(err) || time.Since(stat.ModTime()) > maxCacheAge
}

func getForecast(key string, latitude string, longitude string) (json []byte, err error) {
	res, err := forecast.GetResponse(key, latitude, longitude, "now", "us")
	if err != nil {
		return nil, err
	}

	json, err = ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}

	return json, nil
}

func writeCache(cacheFile string, json []byte) (err error) {
	return ioutil.WriteFile(cacheFile, json, 0644)
}

func formatConditions(condition string) (icon string) {
	icon, ok := conditionIcons[condition]
	if !ok {
		icon = condition
	}
	return
}

func extractConditionFromJSON(jsonBlob []byte) (condition string) {
	f, err := forecast.FromJSON(jsonBlob)
	if err != nil {
		return "❗"
	}

	return f.Currently.Icon
}

func exitWith(message interface{}, status int) {
	fmt.Printf("%v\n", message)
	os.Exit(status)
}

func check(err error) {
	if err != nil {
		exitWith(err, 1)
	}
}
<commit_msg>Use mlbright's forecast<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"

	"github.com/mlbright/forecast/v2"
)

var conditionIcons = map[string]string{
	"clear-day":           "☀",
	"clear-night":         "🌙",
	"cloudy":              "☁",
	"fog":                 "🌁",
	"partly-cloudy-day":   "⛅",
	"partly-cloudy-night": "🌙",
	"rain":                "☔",
	"sleet":               "❄ ☔",
	"snow":                "❄",
	"wind":                "🍃",
}

var maxCacheAge, _ = time.ParseDuration("1h")

func main() {
	coordinates := flag.String("coordinates", "39.95,-75.1667", "the coordinates, expressed as latitude,longitude")
	tmpDir := flag.String("tmpdir", os.TempDir(), "the directory to use to store cached responses")
	key := flag.String("key", os.Getenv("FORECAST_IO_API_KEY"), "your forecast.io API key")

	if *key == "" {
		exitWith("Please provide your forecast.io API key with -key, or set FORECAST_IO_API_KEY", 1)
	}

	flag.Parse()

	coordinateParts := strings.Split(*coordinates, ",")
	var latitude string
	var longitude string

	if len(coordinateParts) != 2 {
		exitWith("You must specify latitude and longitude like so: 39.95,-75.1667", 1)
	} else {
		latitude = coordinateParts[0]
		longitude = coordinateParts[1]
	}

	var cacheFilename = fmt.Sprintf("emoji-weather-%s-%s.json", latitude, longitude)
	var cacheFile = path.Join(*tmpDir, cacheFilename)

	var json []byte
	var err error

	if isCacheStale(cacheFile) {
		json, err = getForecast(*key, latitude, longitude)
		check(err)

		err = writeCache(cacheFile, json)
		check(err)
	} else {
		json, err = ioutil.ReadFile(cacheFile)
		check(err)
	}

	fmt.Println(formatConditions(extractConditionFromJSON(json)))
}

func isCacheStale(cacheFile string) bool {
	stat, err := os.Stat(cacheFile)

	return os.IsNotExist(err) || time.Since(stat.ModTime()) > maxCacheAge
}

func getForecast(key string, latitude string, longitude string) (json []byte, err error) {
	res, err := forecast.GetResponse(key, latitude, longitude, "now", "us")
	if err != nil {
		return nil, err
	}

	json, err = ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}

	return json, nil
}

func writeCache(cacheFile string, json []byte) (err error) {
	return ioutil.WriteFile(cacheFile, json, 0644)
}

func formatConditions(condition string) (icon string) {
	icon, ok := conditionIcons[condition]
	if !ok {
		icon = condition
	}
	return
}

func extractConditionFromJSON(jsonBlob []byte) (condition string) {
	f, err := forecast.FromJSON(jsonBlob)
	if err != nil {
		return "❗"
	}

	return f.Currently.Icon
}

func exitWith(message interface{}, status int) {
	fmt.Printf("%v\n", message)
	os.Exit(status)
}

func check(err error) {
	if err != nil {
		exitWith(err, 1)
	}
}
<|endoftext|>


<commit_before>package typhon

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/golang/protobuf/proto"
	"github.com/monzo/terrors"
)

// A Response is Typhon's wrapper around http.Response, used by both clients and servers.
//
// Note that Typhon makes no guarantees that a Response is safe to access or mutate concurrently. If a single Response
// object is to be used by multiple goroutines concurrently, callers must make sure to properly synchronise accesses.
type Response struct {
	*http.Response
	Error    error
	Request  *Request // The Request that we are responding to
	hijacked bool
}

// Encode serialises the passed object into the body (and sets appropriate headers).
func (r *Response) Encode(v interface{}) {
	if r.Response == nil {
		r.Response = newHTTPResponse(Request{}, http.StatusOK)
	}

	// If we have a proto message of a
	// If we were given an io.ReadCloser or an io.Reader (that is not also a json.Marshaler), use it directly
	switch v := v.(type) {
	case proto.Message, json.Marshaler:
	case io.ReadCloser:
		r.Body = v
		r.ContentLength = -1
		return
	case io.Reader:
		r.Body = ioutil.NopCloser(v)
		r.ContentLength = -1
		return
	}

	// If our request indicates protobuf support and we have a protobuf message
	// then prefer to encode it as a protobuf response body
	acceptsProtobuf := r.Request != nil && strings.Contains(r.Request.Header.Get("Accept"), "application/protobuf")
	if m, ok := v.(proto.Message); ok && acceptsProtobuf {
		r.encodeAsProtobuf(m)
		return
	}

	if err := json.NewEncoder(r).Encode(v); err != nil {
		r.Error = terrors.Wrap(err, nil)
		return
	}
	r.Header.Set("Content-Type", "application/json")
}

func (r *Response) encodeAsProtobuf(m proto.Message) {
	b, err := proto.Marshal(m)
	if err != nil {
		r.Error = terrors.Wrap(err, nil)
		return
	}

	n, err := r.Write(b)
	r.Error = terrors.Wrap(err, nil)
	r.Header.Set("Content-Type", "application/protobuf")
	r.ContentLength = int64(n)
}

// WrapDownstreamErrors is a context key that can be used to enable
// wrapping of downstream response errors on a per-request basis.
//
// This is implemented as a context key to allow us to migrate individual
// services from the old behaviour to the new behaviour without adding a
// dependency on config to Typhon.
type WrapDownstreamErrors struct{}

// Decode de-serialises the body into the passed object.
func (r *Response) Decode(v interface{}) error {
	if r.Error != nil {
		if r.Request != nil && r.Request.Context != nil {
			if s, ok := r.Request.Context.Value(WrapDownstreamErrors{}).(string); ok && s != "" {
				return terrors.NewInternalWithCause(r.Error, "Downstream request error", nil, "downstream")
			}
		}

		return r.Error
	}

	if r.Response == nil {
		r.Error = terrors.InternalService("", "Response has no body", nil)
		return r.Error
	}

	var b []byte
	b, err := r.BodyBytes(true)
	if err != nil {
		r.Error = terrors.WrapWithCode(err, nil, terrors.ErrBadResponse)
		return r.Error
	}

	switch r.Header.Get("Content-Type") {
	case "application/octet-stream", "application/x-google-protobuf", "application/protobuf":
		m, ok := v.(proto.Message)
		if !ok {
			return terrors.InternalService("invalid_type", "could not decode proto message", nil)
		}
		err = proto.Unmarshal(b, m)
	default:
		err = json.Unmarshal(b, v)
	}

	if err != nil {
		r.Error = err
	}
	return err
}

// Write writes the passed bytes to the response's body.
func (r *Response) Write(b []byte) (n int, err error) {
	if r.Response == nil {
		r.Response = newHTTPResponse(Request{}, http.StatusOK)
	}
	switch rc := r.Body.(type) {
	// In the "regular" case, the response body will be a bufCloser; we can write
	case io.Writer:
		n, err = rc.Write(b)
		if err != nil {
			return n, err
		}
	// If a caller manually sets Response.Body, then we may not be able to write to it. In that case, we need to be
	// cleverer.
	default:
		buf := &bufCloser{}
		if rc != nil {
			if _, err := io.Copy(buf, rc); err != nil {
				// This can be quite bad; we have consumed (and possibly lost) some of the original body
				return 0, err
			}
			// rc will never again be accessible: once it's copied it must be closed
			rc.Close()
		}
		r.Body = buf
		n, err = buf.Write(b)
		if err != nil {
			return n, err
		}
	}

	if r.ContentLength >= 0 {
		r.ContentLength += int64(n)
		// If this write pushed the content length above the chunking threshold,
		// set to -1 (unknown) to trigger chunked encoding
		if r.ContentLength >= chunkThreshold {
			r.ContentLength = -1
		}
	}
	return n, nil
}

// BodyBytes fully reads the response body and returns the bytes read. If consume is false, the body is copied into a
// new buffer such that it may be read again.
func (r *Response) BodyBytes(consume bool) ([]byte, error) {
	if consume {
		defer r.Body.Close()
		return ioutil.ReadAll(r.Body)
	}

	switch rc := r.Body.(type) {
	case *bufCloser:
		return rc.Bytes(), nil

	default:
		buf := &bufCloser{}
		r.Body = buf
		rdr := io.TeeReader(rc, buf)
		// rc will never again be accessible: once it's copied it must be closed
		defer rc.Close()
		return ioutil.ReadAll(rdr)
	}
}

// Writer returns a ResponseWriter which can be used to populate the response.
//
// This is useful when you want to use another HTTP library that is used to wrapping net/http directly. For example,
// it allows a Typhon Service to use a http.Handler internally.
func (r *Response) Writer() ResponseWriter {
	if r.Request != nil && r.Request.hijacker != nil {
		return hijackerRw{
			responseWriterWrapper: responseWriterWrapper{
				r: r},
			Hijacker: r.Request.hijacker}
	}
	return responseWriterWrapper{
		r: r}
}

func (r Response) String() string {
	b := new(bytes.Buffer)
	fmt.Fprint(b, "Response(")
	if r.Response != nil {
		fmt.Fprintf(b, "%d", r.StatusCode)
	} else {
		fmt.Fprint(b, "???")
	}
	if r.Error != nil {
		fmt.Fprintf(b, ", error: %v", r.Error)
	}
	fmt.Fprint(b, ")")
	return b.String()
}

func newHTTPResponse(req Request, statusCode int) *http.Response {
	return &http.Response{
		StatusCode:    statusCode,
		Proto:         req.Proto,
		ProtoMajor:    req.ProtoMajor,
		ProtoMinor:    req.ProtoMinor,
		ContentLength: 0,
		Header:        make(http.Header, 5),
		Body:          &bufCloser{}}
}

// NewResponse constructs a Response with status code 200.
func NewResponse(req Request) Response {
	return NewResponseWithCode(req, http.StatusOK)
}

// NewResponseWithCode constructs a Response with the given status code.
func NewResponseWithCode(req Request, statusCode int) Response {
	return Response{
		Request:  &req,
		Error:    nil,
		Response: newHTTPResponse(req, statusCode)}
}
<commit_msg>Remove stray comment<commit_after>package typhon

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/golang/protobuf/proto"
	"github.com/monzo/terrors"
)

// A Response is Typhon's wrapper around http.Response, used by both clients and servers.
//
// Note that Typhon makes no guarantees that a Response is safe to access or mutate concurrently. If a single Response
// object is to be used by multiple goroutines concurrently, callers must make sure to properly synchronise accesses.
type Response struct {
	*http.Response
	Error    error
	Request  *Request // The Request that we are responding to
	hijacked bool
}

// Encode serialises the passed object into the body (and sets appropriate headers).
func (r *Response) Encode(v interface{}) {
	if r.Response == nil {
		r.Response = newHTTPResponse(Request{}, http.StatusOK)
	}

	// If we were given an io.ReadCloser or an io.Reader (that is not also
	// a json.Marshaler or proto.Message), use it directly
	switch v := v.(type) {
	case proto.Message, json.Marshaler:
	case io.ReadCloser:
		r.Body = v
		r.ContentLength = -1
		return
	case io.Reader:
		r.Body = ioutil.NopCloser(v)
		r.ContentLength = -1
		return
	}

	// If our request indicates protobuf support and we have a protobuf message
	// then prefer to encode it as a protobuf response body
	acceptsProtobuf := r.Request != nil && strings.Contains(r.Request.Header.Get("Accept"), "application/protobuf")
	if m, ok := v.(proto.Message); ok && acceptsProtobuf {
		r.encodeAsProtobuf(m)
		return
	}

	if err := json.NewEncoder(r).Encode(v); err != nil {
		r.Error = terrors.Wrap(err, nil)
		return
	}
	r.Header.Set("Content-Type", "application/json")
}

func (r *Response) encodeAsProtobuf(m proto.Message) {
	b, err := proto.Marshal(m)
	if err != nil {
		r.Error = terrors.Wrap(err, nil)
		return
	}

	n, err := r.Write(b)
	r.Error = terrors.Wrap(err, nil)
	r.Header.Set("Content-Type", "application/protobuf")
	r.ContentLength = int64(n)
}

// WrapDownstreamErrors is a context key that can be used to enable
// wrapping of downstream response errors on a per-request basis.
//
// This is implemented as a context key to allow us to migrate individual
// services from the old behaviour to the new behaviour without adding a
// dependency on config to Typhon.
type WrapDownstreamErrors struct{}

// Decode de-serialises the body into the passed object.
func (r *Response) Decode(v interface{}) error {
	if r.Error != nil {
		if r.Request != nil && r.Request.Context != nil {
			if s, ok := r.Request.Context.Value(WrapDownstreamErrors{}).(string); ok && s != "" {
				return terrors.NewInternalWithCause(r.Error, "Downstream request error", nil, "downstream")
			}
		}

		return r.Error
	}

	if r.Response == nil {
		r.Error = terrors.InternalService("", "Response has no body", nil)
		return r.Error
	}

	var b []byte
	b, err := r.BodyBytes(true)
	if err != nil {
		r.Error = terrors.WrapWithCode(err, nil, terrors.ErrBadResponse)
		return r.Error
	}

	switch r.Header.Get("Content-Type") {
	case "application/octet-stream", "application/x-google-protobuf", "application/protobuf":
		m, ok := v.(proto.Message)
		if !ok {
			return terrors.InternalService("invalid_type", "could not decode proto message", nil)
		}
		err = proto.Unmarshal(b, m)
	default:
		err = json.Unmarshal(b, v)
	}

	if err != nil {
		r.Error = err
	}
	return err
}

// Write writes the passed bytes to the response's body.
func (r *Response) Write(b []byte) (n int, err error) {
	if r.Response == nil {
		r.Response = newHTTPResponse(Request{}, http.StatusOK)
	}
	switch rc := r.Body.(type) {
	// In the "regular" case, the response body will be a bufCloser; we can write
	case io.Writer:
		n, err = rc.Write(b)
		if err != nil {
			return n, err
		}
	// If a caller manually sets Response.Body, then we may not be able to write to it. In that case, we need to be
	// cleverer.
	default:
		buf := &bufCloser{}
		if rc != nil {
			if _, err := io.Copy(buf, rc); err != nil {
				// This can be quite bad; we have consumed (and possibly lost) some of the original body
				return 0, err
			}
			// rc will never again be accessible: once it's copied it must be closed
			rc.Close()
		}
		r.Body = buf
		n, err = buf.Write(b)
		if err != nil {
			return n, err
		}
	}

	if r.ContentLength >= 0 {
		r.ContentLength += int64(n)
		// If this write pushed the content length above the chunking threshold,
		// set to -1 (unknown) to trigger chunked encoding
		if r.ContentLength >= chunkThreshold {
			r.ContentLength = -1
		}
	}
	return n, nil
}

// BodyBytes fully reads the response body and returns the bytes read. If consume is false, the body is copied into a
// new buffer such that it may be read again.
func (r *Response) BodyBytes(consume bool) ([]byte, error) {
	if consume {
		defer r.Body.Close()
		return ioutil.ReadAll(r.Body)
	}

	switch rc := r.Body.(type) {
	case *bufCloser:
		return rc.Bytes(), nil

	default:
		buf := &bufCloser{}
		r.Body = buf
		rdr := io.TeeReader(rc, buf)
		// rc will never again be accessible: once it's copied it must be closed
		defer rc.Close()
		return ioutil.ReadAll(rdr)
	}
}

// Writer returns a ResponseWriter which can be used to populate the response.
//
// This is useful when you want to use another HTTP library that is used to wrapping net/http directly. For example,
// it allows a Typhon Service to use a http.Handler internally.
func (r *Response) Writer() ResponseWriter {
	if r.Request != nil && r.Request.hijacker != nil {
		return hijackerRw{
			responseWriterWrapper: responseWriterWrapper{
				r: r},
			Hijacker: r.Request.hijacker}
	}
	return responseWriterWrapper{
		r: r}
}

func (r Response) String() string {
	b := new(bytes.Buffer)
	fmt.Fprint(b, "Response(")
	if r.Response != nil {
		fmt.Fprintf(b, "%d", r.StatusCode)
	} else {
		fmt.Fprint(b, "???")
	}
	if r.Error != nil {
		fmt.Fprintf(b, ", error: %v", r.Error)
	}
	fmt.Fprint(b, ")")
	return b.String()
}

func newHTTPResponse(req Request, statusCode int) *http.Response {
	return &http.Response{
		StatusCode:    statusCode,
		Proto:         req.Proto,
		ProtoMajor:    req.ProtoMajor,
		ProtoMinor:    req.ProtoMinor,
		ContentLength: 0,
		Header:        make(http.Header, 5),
		Body:          &bufCloser{}}
}

// NewResponse constructs a Response with status code 200.
func NewResponse(req Request) Response {
	return NewResponseWithCode(req, http.StatusOK)
}

// NewResponseWithCode constructs a Response with the given status code.
func NewResponseWithCode(req Request, statusCode int) Response {
	return Response{
		Request:  &req,
		Error:    nil,
		Response: newHTTPResponse(req, statusCode)}
}
<|endoftext|>


<commit_before>package engine

// This source file is for the special case of serving a single file.

import (
	"errors"
	"net/http"
	"strconv"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/xyproto/algernon/utils"
	"github.com/xyproto/datablock"
)

const (
	defaultStaticCacheSize = 128 * utils.MiB

	maxAttemptsAtIncreasingPortNumber = 128

	delayBeforeLaunchingBrowser = time.Millisecond * 200
)

// nextPort increases the port number by 1
func nextPort(colonPort string) (string, error) {
	if !strings.HasPrefix(colonPort, ":") {
		return colonPort, errors.New("colonPort does not start with a colon! \"" + colonPort + "\"")
	}
	num, err := strconv.Atoi(colonPort[1:])
	if err != nil {
		return colonPort, errors.New("Could not convert port number to string: \"" + colonPort[1:] + "\"")
	}
	// Increase the port number by 1, add a colon, convert to string and return
	return ":" + strconv.Itoa(num+1), nil
}

// This is a bit hacky, but it's only used when serving a single static file
func (ac *Config) openAfter(wait time.Duration, hostname, colonPort string, https bool, cancelChannel chan bool) {
	// Wait a bit
	time.Sleep(wait)
	select {
	case <-cancelChannel:
		// Got a message on the cancelChannel:
		// don't open the URL with an external application.
		return
	case <-time.After(delayBeforeLaunchingBrowser):
		// Got timeout, assume the port was not busy
		ac.OpenURL(hostname, colonPort, https)
	}
}

// shortInfo outputs a short string about which file is served where
func (ac *Config) shortInfoAndOpen(filename, colonPort string, cancelChannel chan bool) {
	hostname := "localhost"
	if ac.serverHost != "" {
		hostname = ac.serverHost
	}
	log.Info("Serving " + filename + " on http://" + hostname + colonPort)

	if ac.openURLAfterServing {
		go ac.openAfter(delayBeforeLaunchingBrowser, hostname, colonPort, false, cancelChannel)
	}
}

// ServeStaticFile is a convenience function for serving only a single file.
// It can be used as a quick and easy way to view a README.md file.
func (ac *Config) ServeStaticFile(filename, colonPort string) error {
	log.Info("Single file mode. Not using the regular parameters.")

	cancelChannel := make(chan bool, 1)

	ac.shortInfoAndOpen(filename, colonPort, cancelChannel)

	mux := http.NewServeMux()
	// 64 MiB cache, use cache compression, no per-file size limit, use best gzip compression, compress for size not for speed
	ac.cache = datablock.NewFileCache(defaultStaticCacheSize, true, 0, false, 0)
	mux.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		w.Header().Set("Server", ac.versionString)
		ac.FilePage(w, req, filename, ac.defaultLuaDataFilename)
	})
	HTTPserver := ac.NewGracefulServer(mux, false, ac.serverHost+colonPort)

	// Attempt to serve just the single file
	if errServe := HTTPserver.ListenAndServe(); errServe != nil {
		// If it fails, try several times, increasing the port by 1 each time
		for i := 0; i < maxAttemptsAtIncreasingPortNumber; i++ {
			if errServe = HTTPserver.ListenAndServe(); errServe != nil {
				cancelChannel <- true
				if !strings.HasSuffix(errServe.Error(), "already in use") {
					// Not a problem with address already being in use
					ac.fatalExit(errServe)
				}
				log.Warn("Address already in use. Using next port number.")
				if newPort, errNext := nextPort(colonPort); errNext != nil {
					ac.fatalExit(errNext)
				} else {
					colonPort = newPort
				}

				// Make a new cancel channel, and use the new URL
				cancelChannel = make(chan bool, 1)
				ac.shortInfoAndOpen(filename, colonPort, cancelChannel)

				HTTPserver = ac.NewGracefulServer(mux, false, ac.serverHost+colonPort)
			}
		}
		// Several attempts failed
		return errServe
		//ac.fatalExit(errServe)
	}
	return nil
}
<commit_msg>Minor change<commit_after>package engine

// This source file is for the special case of serving a single file.

import (
	"errors"
	"net/http"
	"strconv"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/xyproto/algernon/utils"
	"github.com/xyproto/datablock"
)

const (
	defaultStaticCacheSize            = 128 * utils.MiB
	maxAttemptsAtIncreasingPortNumber = 128
	delayBeforeLaunchingBrowser       = time.Millisecond * 200
)

// nextPort increases the port number by 1
func nextPort(colonPort string) (string, error) {
	if !strings.HasPrefix(colonPort, ":") {
		return colonPort, errors.New("colonPort does not start with a colon! \"" + colonPort + "\"")
	}
	num, err := strconv.Atoi(colonPort[1:])
	if err != nil {
		return colonPort, errors.New("Could not convert port number to string: \"" + colonPort[1:] + "\"")
	}
	// Increase the port number by 1, add a colon, convert to string and return
	return ":" + strconv.Itoa(num+1), nil
}

// This is a bit hacky, but it's only used when serving a single static file
func (ac *Config) openAfter(wait time.Duration, hostname, colonPort string, https bool, cancelChannel chan bool) {
	// Wait a bit
	time.Sleep(wait)
	select {
	case <-cancelChannel:
		// Got a message on the cancelChannel:
		// don't open the URL with an external application.
		return
	case <-time.After(delayBeforeLaunchingBrowser):
		// Got timeout, assume the port was not busy
		ac.OpenURL(hostname, colonPort, https)
	}
}

// shortInfo outputs a short string about which file is served where
func (ac *Config) shortInfoAndOpen(filename, colonPort string, cancelChannel chan bool) {
	hostname := "localhost"
	if ac.serverHost != "" {
		hostname = ac.serverHost
	}
	log.Info("Serving " + filename + " on http://" + hostname + colonPort)

	if ac.openURLAfterServing {
		go ac.openAfter(delayBeforeLaunchingBrowser, hostname, colonPort, false, cancelChannel)
	}
}

// ServeStaticFile is a convenience function for serving only a single file.
// It can be used as a quick and easy way to view a README.md file.
func (ac *Config) ServeStaticFile(filename, colonPort string) error {
	log.Info("Single file mode. Not using the regular parameters.")

	cancelChannel := make(chan bool, 1)

	ac.shortInfoAndOpen(filename, colonPort, cancelChannel)

	mux := http.NewServeMux()
	// 64 MiB cache, use cache compression, no per-file size limit, use best gzip compression, compress for size not for speed
	ac.cache = datablock.NewFileCache(defaultStaticCacheSize, true, 0, false, 0)
	mux.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		w.Header().Set("Server", ac.versionString)
		ac.FilePage(w, req, filename, ac.defaultLuaDataFilename)
	})
	HTTPserver := ac.NewGracefulServer(mux, false, ac.serverHost+colonPort)

	// Attempt to serve just the single file
	if errServe := HTTPserver.ListenAndServe(); errServe != nil {
		// If it fails, try several times, increasing the port by 1 each time
		for i := 0; i < maxAttemptsAtIncreasingPortNumber; i++ {
			if errServe = HTTPserver.ListenAndServe(); errServe != nil {
				cancelChannel <- true
				if !strings.HasSuffix(errServe.Error(), "already in use") {
					// Not a problem with address already being in use
					ac.fatalExit(errServe)
				}
				log.Warn("Address already in use. Using next port number.")
				if newPort, errNext := nextPort(colonPort); errNext != nil {
					ac.fatalExit(errNext)
				} else {
					colonPort = newPort
				}

				// Make a new cancel channel, and use the new URL
				cancelChannel = make(chan bool, 1)
				ac.shortInfoAndOpen(filename, colonPort, cancelChannel)

				HTTPserver = ac.NewGracefulServer(mux, false, ac.serverHost+colonPort)
			}
		}
		// Several attempts failed
		return errServe
		//ac.fatalExit(errServe)
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright (C) 2017 ScyllaDB
// Use of this source code is governed by a ALv2-style
// license that can be found in the LICENSE file.

package qb

import (
	"bytes"
	"fmt"
	"time"
)

// TTL converts duration to format expected in USING TTL clause.
func TTL(d time.Duration) int64 {
	return int64(d.Seconds())
}

// Timestamp converts time to format expected in USING TIMESTAMP clause.
func Timestamp(t time.Time) int64 {
	return t.UnixNano() / 1000
}

type using struct {
	ttl           int64
	ttlName       string
	timestamp     int64
	timestampName string
	timeout       time.Duration
	timeoutName   string

	using bool
}

func (u *using) TTL(d time.Duration) *using {
	u.ttl = TTL(d)
	if u.ttl == 0 {
		u.ttl = -1
	}
	u.timestampName = ""
	return u
}

func (u *using) TTLNamed(name string) *using {
	u.ttl = 0
	u.ttlName = name
	return u
}

func (u *using) Timestamp(t time.Time) *using {
	u.timestamp = Timestamp(t)
	u.timestampName = ""
	return u
}

func (u *using) TimestampNamed(name string) *using {
	u.timestamp = 0
	u.timestampName = name
	return u
}

func (u *using) Timeout(d time.Duration) *using {
	u.timeout = d
	u.timeoutName = ""
	return u
}

func (u *using) TimeoutNamed(name string) *using {
	u.timeout = 0
	u.timeoutName = name
	return u
}

func (u *using) writeCql(cql *bytes.Buffer) (names []string) {
	u.using = false

	if u.ttl != 0 {
		if u.ttl == -1 {
			u.ttl = 0
		}
		u.writePreamble(cql)
		fmt.Fprintf(cql, "TTL %d ", u.ttl)
	} else if u.ttlName != "" {
		u.writePreamble(cql)
		cql.WriteString("TTL ? ")
		names = append(names, u.ttlName)
	}

	if u.timestamp != 0 {
		u.writePreamble(cql)
		fmt.Fprintf(cql, "TIMESTAMP %d ", u.timestamp)
	} else if u.timestampName != "" {
		u.writePreamble(cql)
		cql.WriteString("TIMESTAMP ? ")
		names = append(names, u.timestampName)
	}

	if u.timeout != 0 {
		u.writePreamble(cql)
		fmt.Fprintf(cql, "TIMEOUT %s ", u.timeout)
	} else if u.timeoutName != "" {
		u.writePreamble(cql)
		cql.WriteString("TIMEOUT ? ")
		names = append(names, u.timeoutName)
	}

	return
}

func (u *using) writePreamble(cql *bytes.Buffer) {
	if u.using {
		cql.WriteString("AND ")
	} else {
		cql.WriteString("USING ")
		u.using = true
	}
}
<commit_msg>qb: avoid data race if rendering query builder in different go routines<commit_after>// Copyright (C) 2017 ScyllaDB
// Use of this source code is governed by a ALv2-style
// license that can be found in the LICENSE file.

package qb

import (
	"bytes"
	"fmt"
	"time"
)

// TTL converts duration to format expected in USING TTL clause.
func TTL(d time.Duration) int64 {
	return int64(d.Seconds())
}

// Timestamp converts time to format expected in USING TIMESTAMP clause.
func Timestamp(t time.Time) int64 {
	return t.UnixNano() / 1000
}

type using struct {
	ttl           int64
	ttlName       string
	timestamp     int64
	timestampName string
	timeout       time.Duration
	timeoutName   string
}

func (u *using) TTL(d time.Duration) *using {
	u.ttl = TTL(d)
	if u.ttl == 0 {
		u.ttl = -1
	}
	u.timestampName = ""
	return u
}

func (u *using) TTLNamed(name string) *using {
	u.ttl = 0
	u.ttlName = name
	return u
}

func (u *using) Timestamp(t time.Time) *using {
	u.timestamp = Timestamp(t)
	u.timestampName = ""
	return u
}

func (u *using) TimestampNamed(name string) *using {
	u.timestamp = 0
	u.timestampName = name
	return u
}

func (u *using) Timeout(d time.Duration) *using {
	u.timeout = d
	u.timeoutName = ""
	return u
}

func (u *using) TimeoutNamed(name string) *using {
	u.timeout = 0
	u.timeoutName = name
	return u
}

func (u *using) writeCql(cql *bytes.Buffer) (names []string) {
	writePreamble := u.preambleWriter()

	if u.ttl != 0 {
		if u.ttl == -1 {
			u.ttl = 0
		}
		writePreamble(cql)
		fmt.Fprintf(cql, "TTL %d ", u.ttl)
	} else if u.ttlName != "" {
		writePreamble(cql)
		cql.WriteString("TTL ? ")
		names = append(names, u.ttlName)
	}

	if u.timestamp != 0 {
		writePreamble(cql)
		fmt.Fprintf(cql, "TIMESTAMP %d ", u.timestamp)
	} else if u.timestampName != "" {
		writePreamble(cql)
		cql.WriteString("TIMESTAMP ? ")
		names = append(names, u.timestampName)
	}

	if u.timeout != 0 {
		writePreamble(cql)
		fmt.Fprintf(cql, "TIMEOUT %s ", u.timeout)
	} else if u.timeoutName != "" {
		writePreamble(cql)
		cql.WriteString("TIMEOUT ? ")
		names = append(names, u.timeoutName)
	}

	return
}

func (u *using) preambleWriter() func(cql *bytes.Buffer) {
	var hasPreamble bool
	return func(cql *bytes.Buffer) {
		if hasPreamble {
			cql.WriteString("AND ")
			return
		}
		cql.WriteString("USING ")
		hasPreamble = true
	}
}
<|endoftext|>


<commit_before>package main

// Keep a local cache of DNS packets. Match incoming
// qname,qclass,qtype and return the saved packet.
// On a cache miss consult the nameserver

import (
	"dns"
        "sync"
        "time"
)

// Keep everything in the cache for 60 seconds
const (
        CACHETTL = 1
        _CLASS = 2 << 16

        INSERT = iota
        DELETE
)

var cache Cache

type item struct {
        epoch int64
        msg []byte
}

// Number in the second map denotes the class + type.
func intval(c, t uint16) int {
        return int(c)*_CLASS + int(t)
}

// Mutex entry in the cache, if non-nill take the lock
// Ala Zone in zone.go, but slightly different
type Cache struct {
        data map[string]map[int]*item
        rw  *sync.RWMutex
}

func NewCache() Cache {
        c := new(Cache)
        c.data = make(map[string]map[int]*item)
        c.rw = new(sync.RWMutex)
        return *c
}

// Add an entry to the cache. The old entry (if any) gets overwritten
func (c Cache) add(q *dns.Msg) {
        c.rw.Lock()
        defer c.rw.Unlock()
        qname := q.Question[0].Name
        i := intval(q.Question[0].Qclass, q.Question[0].Qtype)
        if c.data[qname] == nil {
                im := make(map[int]*item)
                c.data[qname] = im
        }
        buf, _ := q.Pack()
        im := c.data[qname]
        im[i] = &item{time.Seconds(), buf}
}

// Lookup an entry in the cache. Returns nil
// when nothing found.
func (c Cache) lookup(q *dns.Msg) []byte {
        // Use the question section for looking up
        c.rw.RLock()
        defer c.rw.RUnlock()
        i := intval(q.Question[0].Qclass, q.Question[0].Qtype)
        if im, ok := c.data[q.Question[0].Name]; ok {
                // we have the name
                if d, ok := im[i]; ok {
                        // We even have the entry, check cache time
                        if time.Seconds() - d.epoch > CACHETTL {
                                // Too olds means we get a new one
                                return nil
                        }
                        e := make([]byte, len(d.msg))
                        copy(e, d.msg)
                        return e
                }
        }
        return nil
}

func checkcache(m *dns.Msg) (o []byte) {
        // Check if we have the packet in Cache
        // if so, return it. Otherwise ask the
        // server, return that answer and put it
        // in the cache.
        o = cache.lookup(m)
        if o != nil {
                // octet 1 and 2 contain the Id, set the one for the current pkt
                o[0] = byte(m.MsgHdr.Id >> 8)
                o[1] = byte(m.MsgHdr.Id)
                return
        }

        println("Cache miss")
        var p *dns.Msg
        for _, c := range qr {
                p, _ = c.Client.Exchange(m, c.Addr)
        }
        cache.add(p)
        o, _ = p.Pack()
        return
}

// Return the configration
func NewFunkenSturm() *FunkenSturm {
	f := new(FunkenSturm)
        f.Funk = make([]*Funk, 1)
	f.Setup = func() bool { cache = NewCache(); return true }
        f.Funk[0] = new(Funk)
        f.Funk[0].Match = func(m *dns.Msg) (*dns.Msg, bool) { return m, true }
	f.Funk[0].Action = checkcache
	return f
}
<commit_msg>Use RawSetId here<commit_after>package main

// Keep a local cache of DNS packets. Match incoming
// qname,qclass,qtype and return the saved packet.
// On a cache miss consult the nameserver

import (
	"dns"
        "sync"
        "time"
)

// Keep everything in the cache for 60 seconds
const (
        CACHETTL = 1
        _CLASS = 2 << 16

        INSERT = iota
        DELETE
)

var cache Cache

type item struct {
        epoch int64
        msg []byte
}

// Number in the second map denotes the class + type.
func intval(c, t uint16) int {
        return int(c)*_CLASS + int(t)
}

// Mutex entry in the cache, if non-nill take the lock
// Ala Zone in zone.go, but slightly different
type Cache struct {
        data map[string]map[int]*item
        rw  *sync.RWMutex
}

func NewCache() Cache {
        c := new(Cache)
        c.data = make(map[string]map[int]*item)
        c.rw = new(sync.RWMutex)
        return *c
}

// Add an entry to the cache. The old entry (if any) gets overwritten
func (c Cache) add(q *dns.Msg) {
        c.rw.Lock()
        defer c.rw.Unlock()
        qname := q.Question[0].Name
        i := intval(q.Question[0].Qclass, q.Question[0].Qtype)
        if c.data[qname] == nil {
                im := make(map[int]*item)
                c.data[qname] = im
        }
        buf, _ := q.Pack()
        im := c.data[qname]
        im[i] = &item{time.Seconds(), buf}
}

// Lookup an entry in the cache. Returns nil
// when nothing found.
func (c Cache) lookup(q *dns.Msg) []byte {
        // Use the question section for looking up
        c.rw.RLock()
        defer c.rw.RUnlock()
        i := intval(q.Question[0].Qclass, q.Question[0].Qtype)
        if im, ok := c.data[q.Question[0].Name]; ok {
                // we have the name
                if d, ok := im[i]; ok {
                        // We even have the entry, check cache time
                        if time.Seconds() - d.epoch > CACHETTL {
                                // Too olds means we get a new one
                                return nil
                        }
                        e := make([]byte, len(d.msg))
                        copy(e, d.msg)
                        return e
                }
        }
        return nil
}

func checkcache(m *dns.Msg) (o []byte) {
        // Check if we have the packet in Cache
        // if so, return it. Otherwise ask the
        // server, return that answer and put it
        // in the cache.
        o = cache.lookup(m)
        if o != nil {
                // octet 1 and 2 contain the Id, set the one for the current pkt
                dns.RawSetId(o, 0, m.MsgHdr.Id)
                return
        }

        println("Cache miss")
        var p *dns.Msg
        for _, c := range qr {
                p, _ = c.Client.Exchange(m, c.Addr)
        }
        cache.add(p)
        o, _ = p.Pack()
        return
}

// Return the configration
func NewFunkenSturm() *FunkenSturm {
	f := new(FunkenSturm)
        f.Funk = make([]*Funk, 1)
	f.Setup = func() bool { cache = NewCache(); return true }
        f.Funk[0] = new(Funk)
        f.Funk[0].Match = func(m *dns.Msg) (*dns.Msg, bool) { return m, true }
	f.Funk[0].Action = checkcache
	return f
}
<|endoftext|>


<commit_before>package eureka

import (
	"log"
	"net/url"
	"github.com/gliderlabs/registrator/bridge"
	eureka "github.com/hudl/fargo"
	"strconv"
)
const DefaultInterval = "10s"

func init() {
	bridge.Register(new(Factory), "eureka")
}

type Factory struct{}

func (f *Factory) New(uri *url.URL) bridge.RegistryAdapter {
	client := eureka.EurekaConnection{}
	if uri.Host != "" {
		client = eureka.NewConn("http://"+uri.Host+uri.Path)
	}else {
		client = eureka.NewConn("http://eureka:8761")
	}

	return &EurekaAdapter{client: client}
}

type EurekaAdapter struct {
	client eureka.EurekaConnection
}

// Ping will try to connect to consul by attempting to retrieve the current leader.
func (r *EurekaAdapter) Ping() error {

	eurekaApps, err := r.client.GetApps()
	if err != nil {
		return err
	}
	log.Println("eureka: current apps ", len(eurekaApps))

	return nil
}

func instanceInformation(service *bridge.Service) *eureka.Instance {

	registration := new(eureka.Instance)

	registration.HostName   = service.IP
	registration.App        = service.Name
	registration.Port       = service.Port
	registration.VipAddress = ShortHandTernary(service.Attrs["eureka_vip"], service.Name)

	if(service.Attrs["eureka_status"] == string(eureka.DOWN)) {
		registration.Status = eureka.DOWN
	} else {
		registration.Status = eureka.UP
	}

	if(service.Attrs["eureka_leaseinfo_renewalintervalinsecs"] != "") {
		v, err := strconv.Atoi(service.Attrs["eureka_leaseinfo_renewalintervalinsecs"])
		if(err != nil) {
			log.Println("eureka: Renewal interval must be valid int", err)
		} else {
			registration.LeaseInfo.RenewalIntervalInSecs = int32(v)
		}
	} else {
		registration.LeaseInfo.RenewalIntervalInSecs = 30
	}

	if(service.Attrs["eureka_leaseinfo_durationinsecs"] != "") {
		v, err := strconv.Atoi(service.Attrs["eureka_leaseinfo_durationinsecs"])
		if(err != nil) {
			log.Println("eureka: Lease duration must be valid int", err)
		} else {
			registration.LeaseInfo.DurationInSecs = int32(v)
		}
	} else {
		registration.LeaseInfo.DurationInSecs = 90
	}

	if service.Attrs["eureka_datacenterinfo_name"] != eureka.MyOwn {
		registration.DataCenterInfo.Name = eureka.Amazon
		registration.DataCenterInfo.Metadata = eureka.AmazonMetadataType {
			PublicHostname: ShortHandTernary(service.Attrs["eureka_datacenterinfo_publichostname"], service.Origin.HostIP),
			PublicIpv4:     ShortHandTernary(service.Attrs["eureka_datacenterinfo_publicipv4"], service.Origin.HostIP),
			LocalHostname:  ShortHandTernary(service.Attrs["eureka_datacenterinfo_localipv4"], service.IP),
			LocalIpv4:      ShortHandTernary(service.Attrs["eureka_datacenterinfo_localhostname"], service.IP),
		}
	} else {
		registration.DataCenterInfo.Name = eureka.MyOwn
	}

	return registration
}

func (r *EurekaAdapter) Register(service *bridge.Service) error {
	registration := instanceInformation(service)
	return r.client.RegisterInstance(registration)
}

func (r *EurekaAdapter) Deregister(service *bridge.Service) error {
	registration := new(eureka.Instance)
	registration.HostName = service.ID
	return r.client.DeregisterInstance(registration)
}

func (r *EurekaAdapter) Refresh(service *bridge.Service) error {
	registration := instanceInformation(service)
	return r.client.ReregisterInstance(registration)
}

func (r *EurekaAdapter) Services() ([]*bridge.Service, error) {
	return []*bridge.Service{}, nil
}

func ShortHandTernary(string1 string, string2 string) string {
	if(string1 != "") {
		return string1
	} else {
		return string2
	}
}<commit_msg>heartbeating instead of re-register<commit_after>package eureka

import (
	"log"
	"net/url"
	"github.com/gliderlabs/registrator/bridge"
	eureka "github.com/hudl/fargo"
	"strconv"
)
const DefaultInterval = "10s"

func init() {
	bridge.Register(new(Factory), "eureka")
}

type Factory struct{}

func (f *Factory) New(uri *url.URL) bridge.RegistryAdapter {
	client := eureka.EurekaConnection{}
	if uri.Host != "" {
		client = eureka.NewConn("http://"+uri.Host+uri.Path)
	}else {
		client = eureka.NewConn("http://eureka:8761")
	}

	return &EurekaAdapter{client: client}
}

type EurekaAdapter struct {
	client eureka.EurekaConnection
}

// Ping will try to connect to consul by attempting to retrieve the current leader.
func (r *EurekaAdapter) Ping() error {

	eurekaApps, err := r.client.GetApps()
	if err != nil {
		return err
	}
	log.Println("eureka: current apps ", len(eurekaApps))

	return nil
}

func instanceInformation(service *bridge.Service) *eureka.Instance {

	registration := new(eureka.Instance)

	registration.HostName   = service.IP
	registration.App        = service.Name
	registration.Port       = service.Port
	registration.VipAddress = ShortHandTernary(service.Attrs["eureka_vip"], service.Name)

	if(service.Attrs["eureka_status"] == string(eureka.DOWN)) {
		registration.Status = eureka.DOWN
	} else {
		registration.Status = eureka.UP
	}

	if(service.Attrs["eureka_leaseinfo_renewalintervalinsecs"] != "") {
		v, err := strconv.Atoi(service.Attrs["eureka_leaseinfo_renewalintervalinsecs"])
		if(err != nil) {
			log.Println("eureka: Renewal interval must be valid int", err)
		} else {
			registration.LeaseInfo.RenewalIntervalInSecs = int32(v)
		}
	} else {
		registration.LeaseInfo.RenewalIntervalInSecs = 30
	}

	if(service.Attrs["eureka_leaseinfo_durationinsecs"] != "") {
		v, err := strconv.Atoi(service.Attrs["eureka_leaseinfo_durationinsecs"])
		if(err != nil) {
			log.Println("eureka: Lease duration must be valid int", err)
		} else {
			registration.LeaseInfo.DurationInSecs = int32(v)
		}
	} else {
		registration.LeaseInfo.DurationInSecs = 90
	}

	if service.Attrs["eureka_datacenterinfo_name"] != eureka.MyOwn {
		registration.DataCenterInfo.Name = eureka.Amazon
		registration.DataCenterInfo.Metadata = eureka.AmazonMetadataType {
			PublicHostname: ShortHandTernary(service.Attrs["eureka_datacenterinfo_publichostname"], service.Origin.HostIP),
			PublicIpv4:     ShortHandTernary(service.Attrs["eureka_datacenterinfo_publicipv4"], service.Origin.HostIP),
			LocalHostname:  ShortHandTernary(service.Attrs["eureka_datacenterinfo_localipv4"], service.IP),
			LocalIpv4:      ShortHandTernary(service.Attrs["eureka_datacenterinfo_localhostname"], service.IP),
		}
	} else {
		registration.DataCenterInfo.Name = eureka.MyOwn
	}

	return registration
}

func (r *EurekaAdapter) Register(service *bridge.Service) error {
	registration := instanceInformation(service)
	return r.client.RegisterInstance(registration)
}

func (r *EurekaAdapter) Deregister(service *bridge.Service) error {
	registration := new(eureka.Instance)
	registration.HostName = service.ID
	return r.client.DeregisterInstance(registration)
}

func (r *EurekaAdapter) Refresh(service *bridge.Service) error {
	registration := instanceInformation(service)
	return r.client.HeartBeatInstance(registration)
}

func (r *EurekaAdapter) Services() ([]*bridge.Service, error) {
	return []*bridge.Service{}, nil
}

func ShortHandTernary(string1 string, string2 string) string {
	if(string1 != "") {
		return string1
	} else {
		return string2
	}
}
<|endoftext|>


<commit_before>package repo

import (
	"encoding/json"
	"io"

	"github.com/hashicorp/go-multierror"
	"github.com/pkg/errors"

	"github.com/FlashbackSRS/flashback-model"
)

// Import imports a .fbb file and stores the content
func Import(user *User, r io.Reader) error {
	pkg := &fb.Package{}
	err := json.NewDecoder(r).Decode(pkg)
	if err != nil {
		return errors.Wrap(err, "Unable to decode JSON")
	}

	udb, err := user.DB()
	if err != nil {
		return errors.Wrap(err, "Unable to connect to User DB")
	}
	bundle := pkg.Bundle
	bdb, err := user.BundleDB(bundle)
	if err != nil {
		return errors.Wrap(err, "Unable to connect to Bundle DB")
	}

	if err := udb.Save(bundle); err != nil {
		return errors.Wrap(err, "Unable to save Bundle to User DB")
	}
	bundle.Rev = nil
	if e := bdb.Save(bundle); e != nil {
		return errors.Wrap(e, "Unable to save Bundle to Bundle DB")
	}

	cardMap := map[string]*fb.Card{}
	for _, c := range pkg.Cards {
		cardMap[c.Identity()] = c
	}

	cards := make([]*fb.Card, 0, len(cardMap))

	for _, d := range pkg.Decks {
		for _, id := range d.Cards.All() {
			c, ok := cardMap[id]
			if !ok {
				return errors.Errorf("Card '%s' listed in deck, but not found in package", id)
			}
			cards = append(cards, c)
		}
	}

	// From this point on, we plow through the errors
	var errs *multierror.Error

	// Themes
	for _, t := range pkg.Themes {
		if err := bdb.Save(t); err != nil {
			errs = multierror.Append(errs, errors.Wrapf(err, "Unable to save Theme %s", t.ID.Identity()))
		}
	}

	// Notes
	for _, n := range pkg.Notes {
		if err := bdb.Save(n); err != nil {
			errs = multierror.Append(errs, errors.Wrapf(err, "Unable to save Note %s", n.ID.Identity()))
		}
	}

	// Decks
	for _, d := range pkg.Decks {
		if err := bdb.Save(d); err != nil {
			errs = multierror.Append(errs, errors.Wrapf(err, "Unable to save Deck %s", d.ID.Identity()))
			continue
		}
	}

	// Cards
	for _, c := range cards {
		if err := udb.Save(c); err != nil {
			errs = multierror.Append(errs, errors.Wrapf(err, "Unable to save Card %s", c.Identity()))
		}
	}
	return errs.ErrorOrNil()
}
<commit_msg>Use bulk updates for imports.<commit_after>package repo

import (
	"encoding/json"
	"io"

	pouchdb "github.com/flimzy/go-pouchdb"
	"github.com/flimzy/log"
	"github.com/hashicorp/go-multierror"
	"github.com/pkg/errors"

	"github.com/FlashbackSRS/flashback-model"
)

// Import imports a .fbb file and stores the content
func Import(user *User, r io.Reader) error {
	pkg := &fb.Package{}
	err := json.NewDecoder(r).Decode(pkg)
	if err != nil {
		return errors.Wrap(err, "Unable to decode JSON")
	}

	udb, err := user.DB()
	if err != nil {
		return errors.Wrap(err, "Unable to connect to User DB")
	}
	bundle := pkg.Bundle
	bdb, err := user.BundleDB(bundle)
	if err != nil {
		return errors.Wrap(err, "Unable to connect to Bundle DB")
	}

	if err := udb.Save(bundle); err != nil {
		return errors.Wrap(err, "Unable to save Bundle to User DB")
	}
	bundle.Rev = nil
	if e := bdb.Save(bundle); e != nil {
		return errors.Wrap(e, "Unable to save Bundle to Bundle DB")
	}

	cardMap := map[string]*fb.Card{}
	for _, c := range pkg.Cards {
		cardMap[c.Identity()] = c
	}

	cards := make([]*fb.Card, 0, len(cardMap))

	for _, d := range pkg.Decks {
		for _, id := range d.Cards.All() {
			c, ok := cardMap[id]
			if !ok {
				return errors.Errorf("Card '%s' listed in deck, but not found in package", id)
			}
			cards = append(cards, c)
		}
	}

	// From this point on, we plow through the errors
	var errs *multierror.Error

	// Themes
	log.Debugln("Saving themes")
	if _, err := bdb.BulkDocs(pkg.Themes, pouchdb.Options{}); err != nil {
		errs = multierror.Append(errs, errors.Wrapf(err, "failure saving themes"))
	}

	// Notes
	log.Debugln("Saving notes")
	if _, err := bdb.BulkDocs(pkg.Notes, pouchdb.Options{}); err != nil {
		errs = multierror.Append(errs, errors.Wrapf(err, "failure saving notes"))
	}

	// Decks
	log.Debugln("Saving decks")
	if _, err := bdb.BulkDocs(pkg.Decks, pouchdb.Options{}); err != nil {
		errs = multierror.Append(errs, errors.Wrapf(err, "failure saving decks"))
	}

	// Cards
	log.Debugln("Saving cards")
	if _, err := udb.BulkDocs(cards, pouchdb.Options{}); err != nil {
		errs = multierror.Append(errs, errors.Wrapf(err, "failure saving cards"))
	}
	return errs.ErrorOrNil()
}
<|endoftext|>


<commit_before>package service

import (
	"errors"
	"fmt"
	"net/http"
	"strings"

	nsq "github.com/bitly/go-nsq"
)

var (
	// ErrEmpty 會在傳入一個空字串時被觸發。
	ErrEmpty = ErrInfo{
		Text:   errors.New("The string is empty."),
		Status: http.StatusBadRequest,
		Code:   "str_empty",
	}
)

type Payload map[string]interface{}

// StringService 是基於字串的服務。
type Service interface {
	Uppercase(string) (string, error)
	Lowercase(string) (string, error)
	Count(string) int
	Test(*nsq.Message)
}

// stringService 概括了字串服務所可用的函式。
type Concrete struct {
	Message *nsq.Producer
}

// ServiceMiddleware 是處理 StringService 的中介層。
type Middleware func(Service) Service

// Uppercase 將傳入的字串轉換為大寫。
func (c Concrete) Uppercase(s string) (string, error) {

	c.Message.Publish("new_user", []byte("test"))

	if s == "" {
		return "", Err{
			Message: ErrEmpty,
		}
	}

	return strings.ToUpper(s), nil
}

// Lowercase 將傳入的字串轉換為小寫。
func (Concrete) Lowercase(s string) (string, error) {
	if s == "" {
		return "", ErrEmpty
	}

	return strings.ToLower(s), nil
}

// Count 計算傳入的字串長度。
func (Concrete) Count(s string) int {
	return len(s)
}

func (Concrete) Test(msg *nsq.Message) {
	fmt.Println(msg)
}
<commit_msg>Removed Lowercase from controller<commit_after>package service

import (
	"errors"
	"fmt"
	"net/http"

	nsq "github.com/bitly/go-nsq"
)

var (
	// ErrEmpty 會在傳入一個空字串時被觸發。
	ErrEmpty = ErrInfo{
		Text:   errors.New("The string is empty."),
		Status: http.StatusBadRequest,
		Code:   "str_empty",
	}
)

// StringService 是基於字串的服務。
type Service interface {
	Uppercase(string) (string, error)
	Count(string) int
	Test(*nsq.Message)
}

// Uppercase 將傳入的字串轉換為大寫。
func (c Concrete) Uppercase(s string) (string, error) {

	//c.Message.Publish("new_user", []byte("test"))

	res, err := c.Model.ToUpper(s)
	if err != nil {
		return "", err
	}

	return res, nil
}

// Count 計算傳入的字串長度。
func (c Concrete) Count(s string) int {
	return c.Model.Count(s)
}

func (Concrete) Test(msg *nsq.Message) {
	fmt.Println(msg)
}
<|endoftext|>