<commit_before>package main

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

import (
	"bytes"
	"compress/gzip"
	"crypto/sha512"
	"database/sql"
	"encoding/base64"
	"fmt"
	"net/http"
	"strings"
	"time"
	"unicode"

	"github.com/apache/incubator-trafficcontrol/traffic_monitor_golang/common/log"
	"github.com/apache/incubator-trafficcontrol/traffic_ops/tocookie"
)

const ServerName = "traffic_ops_golang" + "/" + Version

type AuthRegexHandlerFunc func(w http.ResponseWriter, r *http.Request, params PathParams, user string, privLevel int)

func wrapHeaders(h RegexHandlerFunc) RegexHandlerFunc {
	return func(w http.ResponseWriter, r *http.Request, p PathParams) {
		w.Header().Set("Access-Control-Allow-Credentials", "true")
		w.Header().Set("Access-Control-Allow-Headers", "Origin, X-Requested-With, Content-Type, Accept, Set-Cookie, Cookie")
		w.Header().Set("Access-Control-Allow-Methods", "POST,GET,OPTIONS,PUT,DELETE")
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("X-Server-Name", ServerName)
		iw := &BodyInterceptor{w: w}
		h(iw, r, p)

		sha := sha512.Sum512(iw.Body())
		w.Header().Set("Whole-Content-SHA512", base64.StdEncoding.EncodeToString(sha[:]))

		gzipResponse(w, r, iw.Body())

		iw.RealWrite(iw.Body())

	}
}

func handlerToAuthHandler(h RegexHandlerFunc) AuthRegexHandlerFunc {
	return func(w http.ResponseWriter, r *http.Request, p PathParams, user string, privLevel int) { h(w, r, p) }
}

func wrapAuth(h RegexHandlerFunc, noAuth bool, secret string, privLevelStmt *sql.Stmt, privLevelRequired int) RegexHandlerFunc {
	return wrapAuthWithData(handlerToAuthHandler(h), noAuth, secret, privLevelStmt, privLevelRequired)
}

func wrapAuthWithData(h AuthRegexHandlerFunc, noAuth bool, secret string, privLevelStmt *sql.Stmt, privLevelRequired int) RegexHandlerFunc {
	if noAuth {
		return func(w http.ResponseWriter, r *http.Request, p PathParams) {
			h(w, r, p, "", PrivLevelInvalid)
		}
	}
	return func(w http.ResponseWriter, r *http.Request, p PathParams) {
		// TODO remove, and make username available to wrapLogTime
		start := time.Now()
		iw := &Interceptor{w: w}
		w = iw
		username := "-"
		defer func() {
			log.EventfRaw(`%s - %s [%s] "%v %v HTTP/1.1" %v %v %v "%v"`, r.RemoteAddr, username, time.Now().Format(AccessLogTimeFormat), r.Method, r.URL.Path, iw.code, iw.byteCount, int(time.Now().Sub(start)/time.Millisecond), r.UserAgent())
		}()

		handleUnauthorized := func(reason string) {
			status := http.StatusUnauthorized
			w.WriteHeader(status)
			fmt.Fprintf(w, http.StatusText(status))
			log.Infof("%v %v %v %v returned unauthorized: %v\n", r.RemoteAddr, r.Method, r.URL.Path, username, reason)
		}

		cookie, err := r.Cookie(tocookie.Name)
		if err != nil {
			handleUnauthorized("error getting cookie: " + err.Error())
			return
		}

		if cookie == nil {
			handleUnauthorized("no auth cookie")
			return
		}

		oldCookie, err := tocookie.Parse(secret, cookie.Value)
		if err != nil {
			handleUnauthorized("cookie error: " + err.Error())
			return
		}

		username = oldCookie.AuthData
		privLevel := PrivLevel(privLevelStmt, username)
		if privLevel < privLevelRequired {
			handleUnauthorized("insufficient privileges")
			return
		}

		newCookieVal := tocookie.Refresh(oldCookie, secret)
		http.SetCookie(w, &http.Cookie{Name: tocookie.Name, Value: newCookieVal, Path: "/", HttpOnly: true})

		h(w, r, p, username, privLevel)
	}
}

const AccessLogTimeFormat = "02/Jan/2006:15:04:05 -0700"

func wrapAccessLog(secret string, h http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		iw := &Interceptor{w: w}
		user := "-"
		cookie, err := r.Cookie(tocookie.Name)
		if err == nil && cookie != nil {
			cookie, err := tocookie.Parse(secret, cookie.Value)
			if err == nil {
				user = cookie.AuthData
			}
		}
		start := time.Now()
		defer func() {
			log.EventfRaw(`%s - %s [%s] "%v %v HTTP/1.1" %v %v %v "%v"`, r.RemoteAddr, user, time.Now().Format(AccessLogTimeFormat), r.Method, r.URL.Path, iw.code, iw.byteCount, int(time.Now().Sub(start)/time.Millisecond), r.UserAgent())
		}()
		h.ServeHTTP(iw, r)
	}
}

// gzipResponse takes a function which cannot error and returns only bytes, and wraps it as a http.HandlerFunc. The errContext is logged if the write fails, and should be enough information to trace the problem (function name, endpoint, request parameters, etc).
func gzipResponse(w http.ResponseWriter, r *http.Request, bytes []byte) {

	bytes, err := gzipIfAccepts(r, w, bytes)
	if err != nil {
		log.Errorf("gzipping request '%v': %v\n", r.URL.EscapedPath(), err)
		code := http.StatusInternalServerError
		w.WriteHeader(code)
		if _, err := w.Write([]byte(http.StatusText(code))); err != nil {
			log.Warnf("received error writing data request %v: %v\n", r.URL.EscapedPath(), err)
		}
		return
	}

	w.Write(bytes)
}

// wrapBytes takes a function which cannot error and returns only bytes, and wraps it as a http.HandlerFunc. The errContext is logged if the write fails, and should be enough information to trace the problem (function name, endpoint, request parameters, etc).
//TODO: drichardson - refactor these to a generic area
func wrapBytes(f func() []byte, contentType string) RegexHandlerFunc {
	return func(w http.ResponseWriter, r *http.Request, p PathParams) {
		bytes := f()
		bytes, err := gzipIfAccepts(r, w, bytes)
		if err != nil {
			log.Errorf("gzipping request '%v': %v\n", r.URL.EscapedPath(), err)
			code := http.StatusInternalServerError
			w.WriteHeader(code)
			if _, err := w.Write([]byte(http.StatusText(code))); err != nil {
				log.Warnf("received error writing data request %v: %v\n", r.URL.EscapedPath(), err)
			}
			return
		}

		w.Header().Set("Content-Type", contentType)
		log.Write(w, bytes, r.URL.EscapedPath())
	}
}

// gzipIfAccepts gzips the given bytes, writes a `Content-Encoding: gzip` header to the given writer, and returns the gzipped bytes, if the Request supports GZip (has an Accept-Encoding header). Else, returns the bytes unmodified. Note the given bytes are NOT written to the given writer. It is assumed the bytes may need to pass thru other middleware before being written.
//TODO: drichardson - refactor these to a generic area
func gzipIfAccepts(r *http.Request, w http.ResponseWriter, b []byte) ([]byte, error) {
	// TODO this could be made more efficient by wrapping ResponseWriter with the GzipWriter, and letting callers writer directly to it - but then we'd have to deal with Closing the gzip.Writer.
	if len(b) == 0 || !acceptsGzip(r) {
		return b, nil
	}
	w.Header().Set("Content-Encoding", "gzip")

	buf := bytes.Buffer{}
	zw := gzip.NewWriter(&buf)

	if _, err := zw.Write(b); err != nil {
		return nil, fmt.Errorf("gzipping bytes: %v")
	}

	if err := zw.Close(); err != nil {
		return nil, fmt.Errorf("closing gzip writer: %v")
	}

	return buf.Bytes(), nil
}

func acceptsGzip(r *http.Request) bool {
	encodingHeaders := r.Header["Accept-Encoding"] // headers are case-insensitive, but Go promises to Canonical-Case requests
	for _, encodingHeader := range encodingHeaders {
		encodingHeader = stripAllWhitespace(encodingHeader)
		encodings := strings.Split(encodingHeader, ",")
		for _, encoding := range encodings {
			if strings.ToLower(encoding) == "gzip" { // encoding is case-insensitive, per the RFC
				return true
			}
		}
	}
	return false
}

func stripAllWhitespace(s string) string {
	return strings.Map(func(r rune) rune {
		if unicode.IsSpace(r) {
			return -1
		}
		return r
	}, s)
}

type Interceptor struct {
	w         http.ResponseWriter
	code      int
	byteCount int
}

func (i *Interceptor) WriteHeader(rc int) {
	i.w.WriteHeader(rc)
	i.code = rc
}

func (i *Interceptor) Write(b []byte) (int, error) {
	wi, werr := i.w.Write(b)
	i.byteCount += wi
	if i.code == 0 {
		i.code = 200
	}
	return wi, werr
}

func (i *Interceptor) Header() http.Header {
	return i.w.Header()
}

// BodyInterceptor fulfills the Writer interface, but records the body and doesn't actually write. This allows performing operations on the entire body written by a handler, for example, compressing or hashing. To actually write, call `RealWrite()`. Note this means `len(b)` and `nil` are always returned by `Write()`, any real write errors will be returned by `RealWrite()`.
type BodyInterceptor struct {
	w    http.ResponseWriter
	body []byte
}

func (i *BodyInterceptor) WriteHeader(rc int) {
	i.w.WriteHeader(rc)
}
func (i *BodyInterceptor) Write(b []byte) (int, error) {
	i.body = append(i.body, b...)
	return len(b), nil
}
func (i *BodyInterceptor) Header() http.Header {
	return i.w.Header()
}
func (i *BodyInterceptor) RealWrite(b []byte) (int, error) {
	wi, werr := i.w.Write(i.body)
	return wi, werr
}
func (i *BodyInterceptor) Body() []byte {
	return i.body
}
<commit_msg>added gzip check to prevent double response<commit_after>package main

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

import (
	"bytes"
	"compress/gzip"
	"crypto/sha512"
	"database/sql"
	"encoding/base64"
	"fmt"
	"net/http"
	"strings"
	"time"
	"unicode"

	"github.com/apache/incubator-trafficcontrol/traffic_monitor_golang/common/log"
	"github.com/apache/incubator-trafficcontrol/traffic_ops/tocookie"
)

const ServerName = "traffic_ops_golang" + "/" + Version

type AuthRegexHandlerFunc func(w http.ResponseWriter, r *http.Request, params PathParams, user string, privLevel int)

func wrapHeaders(h RegexHandlerFunc) RegexHandlerFunc {
	return func(w http.ResponseWriter, r *http.Request, p PathParams) {
		w.Header().Set("Access-Control-Allow-Credentials", "true")
		w.Header().Set("Access-Control-Allow-Headers", "Origin, X-Requested-With, Content-Type, Accept, Set-Cookie, Cookie")
		w.Header().Set("Access-Control-Allow-Methods", "POST,GET,OPTIONS,PUT,DELETE")
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("X-Server-Name", ServerName)
		iw := &BodyInterceptor{w: w}
		h(iw, r, p)

		sha := sha512.Sum512(iw.Body())
		w.Header().Set("Whole-Content-SHA512", base64.StdEncoding.EncodeToString(sha[:]))

		if acceptsGzip(r) {
			gzipResponse(w, r, iw.Body())
		} else {
			iw.RealWrite(iw.Body())
		}

	}
}

func handlerToAuthHandler(h RegexHandlerFunc) AuthRegexHandlerFunc {
	return func(w http.ResponseWriter, r *http.Request, p PathParams, user string, privLevel int) { h(w, r, p) }
}

func wrapAuth(h RegexHandlerFunc, noAuth bool, secret string, privLevelStmt *sql.Stmt, privLevelRequired int) RegexHandlerFunc {
	return wrapAuthWithData(handlerToAuthHandler(h), noAuth, secret, privLevelStmt, privLevelRequired)
}

func wrapAuthWithData(h AuthRegexHandlerFunc, noAuth bool, secret string, privLevelStmt *sql.Stmt, privLevelRequired int) RegexHandlerFunc {
	if noAuth {
		return func(w http.ResponseWriter, r *http.Request, p PathParams) {
			h(w, r, p, "", PrivLevelInvalid)
		}
	}
	return func(w http.ResponseWriter, r *http.Request, p PathParams) {
		// TODO remove, and make username available to wrapLogTime
		start := time.Now()
		iw := &Interceptor{w: w}
		w = iw
		username := "-"
		defer func() {
			log.EventfRaw(`%s - %s [%s] "%v %v HTTP/1.1" %v %v %v "%v"`, r.RemoteAddr, username, time.Now().Format(AccessLogTimeFormat), r.Method, r.URL.Path, iw.code, iw.byteCount, int(time.Now().Sub(start)/time.Millisecond), r.UserAgent())
		}()

		handleUnauthorized := func(reason string) {
			status := http.StatusUnauthorized
			w.WriteHeader(status)
			fmt.Fprintf(w, http.StatusText(status))
			log.Infof("%v %v %v %v returned unauthorized: %v\n", r.RemoteAddr, r.Method, r.URL.Path, username, reason)
		}

		cookie, err := r.Cookie(tocookie.Name)
		if err != nil {
			handleUnauthorized("error getting cookie: " + err.Error())
			return
		}

		if cookie == nil {
			handleUnauthorized("no auth cookie")
			return
		}

		oldCookie, err := tocookie.Parse(secret, cookie.Value)
		if err != nil {
			handleUnauthorized("cookie error: " + err.Error())
			return
		}

		username = oldCookie.AuthData
		privLevel := PrivLevel(privLevelStmt, username)
		if privLevel < privLevelRequired {
			handleUnauthorized("insufficient privileges")
			return
		}

		newCookieVal := tocookie.Refresh(oldCookie, secret)
		http.SetCookie(w, &http.Cookie{Name: tocookie.Name, Value: newCookieVal, Path: "/", HttpOnly: true})

		h(w, r, p, username, privLevel)
	}
}

const AccessLogTimeFormat = "02/Jan/2006:15:04:05 -0700"

func wrapAccessLog(secret string, h http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		iw := &Interceptor{w: w}
		user := "-"
		cookie, err := r.Cookie(tocookie.Name)
		if err == nil && cookie != nil {
			cookie, err := tocookie.Parse(secret, cookie.Value)
			if err == nil {
				user = cookie.AuthData
			}
		}
		start := time.Now()
		defer func() {
			log.EventfRaw(`%s - %s [%s] "%v %v HTTP/1.1" %v %v %v "%v"`, r.RemoteAddr, user, time.Now().Format(AccessLogTimeFormat), r.Method, r.URL.Path, iw.code, iw.byteCount, int(time.Now().Sub(start)/time.Millisecond), r.UserAgent())
		}()
		h.ServeHTTP(iw, r)
	}
}

// gzipResponse takes a function which cannot error and returns only bytes, and wraps it as a http.HandlerFunc. The errContext is logged if the write fails, and should be enough information to trace the problem (function name, endpoint, request parameters, etc).
func gzipResponse(w http.ResponseWriter, r *http.Request, bytes []byte) {

	bytes, err := gzipIfAccepts(r, w, bytes)
	if err != nil {
		log.Errorf("gzipping request '%v': %v\n", r.URL.EscapedPath(), err)
		code := http.StatusInternalServerError
		w.WriteHeader(code)
		if _, err := w.Write([]byte(http.StatusText(code))); err != nil {
			log.Warnf("received error writing data request %v: %v\n", r.URL.EscapedPath(), err)
		}
		return
	}

	w.Write(bytes)
}

// wrapBytes takes a function which cannot error and returns only bytes, and wraps it as a http.HandlerFunc. The errContext is logged if the write fails, and should be enough information to trace the problem (function name, endpoint, request parameters, etc).
//TODO: drichardson - refactor these to a generic area
func wrapBytes(f func() []byte, contentType string) RegexHandlerFunc {
	return func(w http.ResponseWriter, r *http.Request, p PathParams) {
		bytes := f()
		bytes, err := gzipIfAccepts(r, w, bytes)
		if err != nil {
			log.Errorf("gzipping request '%v': %v\n", r.URL.EscapedPath(), err)
			code := http.StatusInternalServerError
			w.WriteHeader(code)
			if _, err := w.Write([]byte(http.StatusText(code))); err != nil {
				log.Warnf("received error writing data request %v: %v\n", r.URL.EscapedPath(), err)
			}
			return
		}

		w.Header().Set("Content-Type", contentType)
		log.Write(w, bytes, r.URL.EscapedPath())
	}
}

// gzipIfAccepts gzips the given bytes, writes a `Content-Encoding: gzip` header to the given writer, and returns the gzipped bytes, if the Request supports GZip (has an Accept-Encoding header). Else, returns the bytes unmodified. Note the given bytes are NOT written to the given writer. It is assumed the bytes may need to pass thru other middleware before being written.
//TODO: drichardson - refactor these to a generic area
func gzipIfAccepts(r *http.Request, w http.ResponseWriter, b []byte) ([]byte, error) {
	// TODO this could be made more efficient by wrapping ResponseWriter with the GzipWriter, and letting callers writer directly to it - but then we'd have to deal with Closing the gzip.Writer.
	if len(b) == 0 || !acceptsGzip(r) {
		return b, nil
	}
	w.Header().Set("Content-Encoding", "gzip")

	buf := bytes.Buffer{}
	zw := gzip.NewWriter(&buf)

	if _, err := zw.Write(b); err != nil {
		return nil, fmt.Errorf("gzipping bytes: %v")
	}

	if err := zw.Close(); err != nil {
		return nil, fmt.Errorf("closing gzip writer: %v")
	}

	return buf.Bytes(), nil
}

func acceptsGzip(r *http.Request) bool {
	encodingHeaders := r.Header["Accept-Encoding"] // headers are case-insensitive, but Go promises to Canonical-Case requests
	for _, encodingHeader := range encodingHeaders {
		encodingHeader = stripAllWhitespace(encodingHeader)
		encodings := strings.Split(encodingHeader, ",")
		for _, encoding := range encodings {
			if strings.ToLower(encoding) == "gzip" { // encoding is case-insensitive, per the RFC
				return true
			}
		}
	}
	return false
}

func stripAllWhitespace(s string) string {
	return strings.Map(func(r rune) rune {
		if unicode.IsSpace(r) {
			return -1
		}
		return r
	}, s)
}

type Interceptor struct {
	w         http.ResponseWriter
	code      int
	byteCount int
}

func (i *Interceptor) WriteHeader(rc int) {
	i.w.WriteHeader(rc)
	i.code = rc
}

func (i *Interceptor) Write(b []byte) (int, error) {
	wi, werr := i.w.Write(b)
	i.byteCount += wi
	if i.code == 0 {
		i.code = 200
	}
	return wi, werr
}

func (i *Interceptor) Header() http.Header {
	return i.w.Header()
}

// BodyInterceptor fulfills the Writer interface, but records the body and doesn't actually write. This allows performing operations on the entire body written by a handler, for example, compressing or hashing. To actually write, call `RealWrite()`. Note this means `len(b)` and `nil` are always returned by `Write()`, any real write errors will be returned by `RealWrite()`.
type BodyInterceptor struct {
	w    http.ResponseWriter
	body []byte
}

func (i *BodyInterceptor) WriteHeader(rc int) {
	i.w.WriteHeader(rc)
}
func (i *BodyInterceptor) Write(b []byte) (int, error) {
	i.body = append(i.body, b...)
	return len(b), nil
}
func (i *BodyInterceptor) Header() http.Header {
	return i.w.Header()
}
func (i *BodyInterceptor) RealWrite(b []byte) (int, error) {
	wi, werr := i.w.Write(i.body)
	return wi, werr
}
func (i *BodyInterceptor) Body() []byte {
	return i.body
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Modified 2016 by Steve Manuel, Boss Sauce Creative, LLC
// All modifications are relicensed under the same BSD license
// found in the LICENSE file.

// Generate a self-signed X.509 certificate for a TLS server. Outputs to
// 'devcerts/cert.pem' and 'devcerts/key.pem' and will overwrite existing files.

package tls

import (
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"log"
	"math/big"
	"net"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/ponzu-cms/ponzu/system/db"
)

func publicKey(priv interface{}) interface{} {
	switch k := priv.(type) {
	case *rsa.PrivateKey:
		return &k.PublicKey
	case *ecdsa.PrivateKey:
		return &k.PublicKey
	default:
		return nil
	}
}

func pemBlockForKey(priv interface{}) *pem.Block {
	switch k := priv.(type) {
	case *rsa.PrivateKey:
		return &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(k)}
	case *ecdsa.PrivateKey:
		b, err := x509.MarshalECPrivateKey(k)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unable to marshal ECDSA private key: %v", err)
			os.Exit(2)
		}
		return &pem.Block{Type: "EC PRIVATE KEY", Bytes: b}
	default:
		return nil
	}
}

func setupDev() {
	var priv interface{}
	var err error

	priv, err = rsa.GenerateKey(rand.Reader, 2048)

	if err != nil {
		log.Fatalf("failed to generate private key: %s", err)
	}

	notBefore := time.Now()
	notAfter := notBefore.Add(time.Hour * 24 * 30) // valid for 30 days

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		log.Fatalf("failed to generate serial number: %s", err)
	}

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: []string{"Ponzu Dev Server"},
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	host := db.ConfigCache("domain")
	if host == "" {
		host = "localhost, 0.0.0.0"
	}
	hosts := strings.Split(host, ",")
	for _, h := range hosts {
		if ip := net.ParseIP(h); ip != nil {
			template.IPAddresses = append(template.IPAddresses, ip)
		} else {
			template.DNSNames = append(template.DNSNames, h)
		}
	}

	// make all certs CA
	// template.IsCA = true
	// template.KeyUsage |= x509.KeyUsageCertSign

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, publicKey(priv), priv)
	if err != nil {
		log.Fatalln("Failed to create certificate:", err)
	}

	// overwrite/create directory for devcerts
	pwd, err := os.Getwd()
	if err != nil {
		log.Fatalln("Couldn't find working directory to locate or save dev certificates:", err)
	}

	vendorTLSPath := filepath.Join(pwd, "cmd", "ponzu", "vendor", "github.com", "ponzu-cms", "ponzu", "system", "tls")
	devcertsPath := filepath.Join(vendorTLSPath, "devcerts")
	fmt.Println(devcertsPath)

	err = os.Mkdir(devcertsPath, os.ModePerm|os.ModePerm)
	if err != nil {
		log.Fatalln("Failed to create directory to locate or save dev certificates:", err)
	}

	certOut, err := os.Create(filepath.Join(devcertsPath, "cert.pem"))
	if err != nil {
		log.Fatalln("Failed to open devcerts/cert.pem for writing:", err)
	}
	pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	certOut.Close()

	keyOut, err := os.OpenFile(filepath.Join(devcertsPath, "key.pem"), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		log.Fatalln("Failed to open devcerts/key.pem for writing:", err)
		return
	}
	pem.Encode(keyOut, pemBlockForKey(priv))
	keyOut.Close()
}
<commit_msg>add additional cert usage back for testing<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Modified 2016 by Steve Manuel, Boss Sauce Creative, LLC
// All modifications are relicensed under the same BSD license
// found in the LICENSE file.

// Generate a self-signed X.509 certificate for a TLS server. Outputs to
// 'devcerts/cert.pem' and 'devcerts/key.pem' and will overwrite existing files.

package tls

import (
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"log"
	"math/big"
	"net"
	"os"
	"path/filepath"
	"time"

	"github.com/ponzu-cms/ponzu/system/db"
)

func publicKey(priv interface{}) interface{} {
	switch k := priv.(type) {
	case *rsa.PrivateKey:
		return &k.PublicKey
	case *ecdsa.PrivateKey:
		return &k.PublicKey
	default:
		return nil
	}
}

func pemBlockForKey(priv interface{}) *pem.Block {
	switch k := priv.(type) {
	case *rsa.PrivateKey:
		return &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(k)}
	case *ecdsa.PrivateKey:
		b, err := x509.MarshalECPrivateKey(k)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unable to marshal ECDSA private key: %v", err)
			os.Exit(2)
		}
		return &pem.Block{Type: "EC PRIVATE KEY", Bytes: b}
	default:
		return nil
	}
}

func setupDev() {
	var priv interface{}
	var err error

	priv, err = rsa.GenerateKey(rand.Reader, 2048)

	if err != nil {
		log.Fatalf("failed to generate private key: %s", err)
	}

	notBefore := time.Now()
	notAfter := notBefore.Add(time.Hour * 24 * 30) // valid for 30 days

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		log.Fatalf("failed to generate serial number: %s", err)
	}

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: []string{"Ponzu Dev Server"},
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	hosts := []string{"localhost", "0.0.0.0"}
	domain := db.ConfigCache("domain")
	if domain != "" {
		hosts = append(hosts, domain)
	}

	for _, h := range hosts {
		if ip := net.ParseIP(h); ip != nil {
			template.IPAddresses = append(template.IPAddresses, ip)
		} else {
			template.DNSNames = append(template.DNSNames, h)
		}
	}

	hosts = []string{"localhost", "0.0.0.0"}
	// make all certs CA
	// template.IsCA = true
	template.KeyUsage |= x509.KeyUsageCertSign

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, publicKey(priv), priv)
	if err != nil {
		log.Fatalln("Failed to create certificate:", err)
	}

	// overwrite/create directory for devcerts
	pwd, err := os.Getwd()
	if err != nil {
		log.Fatalln("Couldn't find working directory to locate or save dev certificates:", err)
	}

	vendorTLSPath := filepath.Join(pwd, "cmd", "ponzu", "vendor", "github.com", "ponzu-cms", "ponzu", "system", "tls")
	devcertsPath := filepath.Join(vendorTLSPath, "devcerts")
	fmt.Println(devcertsPath)

	err = os.Mkdir(devcertsPath, os.ModePerm|os.ModePerm)
	if err != nil {
		log.Fatalln("Failed to create directory to locate or save dev certificates:", err)
	}

	certOut, err := os.Create(filepath.Join(devcertsPath, "cert.pem"))
	if err != nil {
		log.Fatalln("Failed to open devcerts/cert.pem for writing:", err)
	}
	pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	certOut.Close()

	keyOut, err := os.OpenFile(filepath.Join(devcertsPath, "key.pem"), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		log.Fatalln("Failed to open devcerts/key.pem for writing:", err)
		return
	}
	pem.Encode(keyOut, pemBlockForKey(priv))
	keyOut.Close()
}
<|endoftext|>


<commit_before>package tasker_test

import (
	"fmt"
	"strings"

	"github.com/krasoffski/gomill/tasker"
)

type upperTask struct {
	in, out string
}

// Process processes and fills required fields of upperTask.
func (u *upperTask) Process() {
	u.out = strings.ToUpper(u.in)
}

// Output prints out Task result to standart output.
func (u *upperTask) Output() {
	fmt.Printf("%s\n", u.out)
}

type taskBuilder struct {
	words   []string
	bufSize int
}

func (tb *taskBuilder) BufSize() int {
	return tb.bufSize
}

func (tb *taskBuilder) Create(s string) tasker.Task {
	return &upperTask{in: s}
}

// Items prepares strings from slice and send them to chan.
func (tb *taskBuilder) Items() <-chan string {
	ch := make(chan string)
	go func() {
		defer close(ch)
		for _, word := range tb.words {
			if word == "" {
				continue
			}
			ch <- word
		}
	}()
	return ch
}

// Run performs all tasks using provided number of goroutines and buffer size.
func (tb *taskBuilder) Run(workers int) {
	tasker.Run(tb, workers)
}

// New creates and initializes new task builder.
func New(input []string, bufSize int) tasker.Builder {
	return &taskBuilder{words: input, bufSize: bufSize}
}

func Example() {
	input := []string{"apple", "orange", "", "cherry"}
	u := New(input, 10)
	// NOTE: the output order might be different due to async processing.
	u.Run(5)
	// Output: APPLE
	// ORANGE
	// CHERRY
}
<commit_msg>playing with example for runner<commit_after>package tasker_test

import (
	"fmt"
	"strings"

	"github.com/krasoffski/gomill/tasker"
)

type upperTask struct {
	in, out string
}

// Process processes and fills required fields of upperTask.
func (u *upperTask) Process() {
	u.out = strings.ToUpper(u.in)
}

// Output prints out Task result to standart output.
func (u *upperTask) Output() {
	fmt.Printf("%s ", u.out)
}

type taskBuilder struct {
	words   []string
	bufSize int
}

func (tb *taskBuilder) BufSize() int {
	return tb.bufSize
}

func (tb *taskBuilder) Create(s string) tasker.Task {
	return &upperTask{in: s}
}

// Items prepares strings from slice and send them to chan.
func (tb *taskBuilder) Items() <-chan string {
	ch := make(chan string)
	go func() {
		defer close(ch)
		for _, word := range tb.words {
			if word == "" {
				continue
			}
			ch <- word
		}
	}()
	return ch
}

// Run performs all tasks using provided number of goroutines and buffer size.
func (tb *taskBuilder) Run(workers int) {
	tasker.Run(tb, workers)
}

// New creates and initializes new task builder.
func New(input []string, bufSize int) tasker.Builder {
	return &taskBuilder{words: input, bufSize: bufSize}
}

func Example() {
	input := []string{"apple", "orange", "", "cherry"}
	u := New(input, 10)
	// NOTE: the output order might be different due to async processing.
	u.Run(5)
	// Output: APPLE ORANGE CHERRY
}
<|endoftext|>


<commit_before>package of10

import (
	"net"

	. "github.com/oshothebig/goflow/openflow"
)

type FeaturesRequest struct {
	Header
}

type FeaturesReply struct {
	Header
	DatapathId   DatapathId
	Buffers      uint32
	Tables       uint8
	pad          [3]uint8
	Capabilities Capability
	Actions      ActionType
	Ports        []PhysicalPort
}

type DatapathId uint64
type Capability uint32

const (
	OFPC_FLOW_STATS Capability = 1 << iota
	OFPC_TABLE_STATS
	OFPC_PORT_STATS
	OFPC_STP
	OFPC_RESERVED
	OFPC_IP_REASM
	OFPC_QUEUE_STATS
	OFPC_ARP_MATCH_IP
)

var Capabilities = struct {
	FlowStats    Capability
	TableStats   Capability
	PortStats    Capability
	Stp          Capability
	Reserved     Capability
	IpReassemble Capability
	QueueStats   Capability
	ArpMatchIp   Capability
}{
	OFPC_FLOW_STATS,
	OFPC_TABLE_STATS,
	OFPC_PORT_STATS,
	OFPC_STP,
	OFPC_RESERVED,
	OFPC_IP_REASM,
	OFPC_QUEUE_STATS,
	OFPC_ARP_MATCH_IP,
}

type SwitchConfig struct {
	Header
	Flags          ConfigFlag
	MissSendLength uint16
}

type ConfigFlag uint16

const (
	OFPC_FRAG_NORMAL ConfigFlag = iota
	OFPC_FRAG_DROP
	OFPC_FRAG_REASM
	OFPC_FRAG_MASK
)

var ConfigFlags = struct {
	FragmentNormal     ConfigFlag
	FragmentDrop       ConfigFlag
	FragmentReassemble ConfigFlag
	FragmentMask       ConfigFlag
}{
	OFPC_FRAG_NORMAL,
	OFPC_FRAG_DROP,
	OFPC_FRAG_REASM,
	OFPC_FRAG_MASK,
}

type FlowMod struct {
	Header
	Match       Match
	Cookie      Cookie
	Command     FlowModCommand
	IdleTimeout uint16
	HardTimeout uint16
	Priority    uint16
	BufferId    BufferId
	OutPort     PortNumber
	Flags       FlowModFlag
	Actions     []Action
}

type Cookie uint64
type FlowModCommand uint16
type BufferId uint32
type FlowModFlag uint16

const (
	OFPFC_ADD FlowModCommand = iota
	OFPFC_MODIFY
	OFPFC_MODIFY_STRICT
	OFPFC_DELETE
	OFPFC_DELETE_STRICT
)

var FlowModCommands = struct {
	Add          FlowModCommand
	Modify       FlowModCommand
	ModifyStrict FlowModCommand
	Delete       FlowModCommand
	DeleteStrict FlowModCommand
}{
	OFPFC_ADD,
	OFPFC_MODIFY,
	OFPFC_MODIFY_STRICT,
	OFPFC_DELETE,
	OFPFC_DELETE_STRICT,
}

const (
	OFPFF_SEND_FLOW_REM FlowModFlag = 1 << iota
	OFPFF_CHECK_OVERLAP
	OFPFF_EMERG
)

var FlowModFlags = struct {
	SendFlowRemoved FlowModFlag
	CheckOverlap    FlowModFlag
	Emergency       FlowModFlag
}{
	OFPFF_SEND_FLOW_REM,
	OFPFF_CHECK_OVERLAP,
	OFPFF_EMERG,
}

type PortMod struct {
	Header
	PortNumber      PortNumber
	HardwareAddress net.HardwareAddr
	Config          PortConfig
	Mask            PortConfig
	Advertise       PortFeature
	pad             [4]uint8
}

type QueueGetConfigRequest struct {
	Header
	Port PortNumber
	pad  [2]uint8
}

type QueueGetConfigReply struct {
	Header
	Port   PortNumber
	pad    [6]uint8
	Queues []PacketQueue
}

type StatsRequest struct {
	Header
	Type  StatsType
	Flags uint16
	Body  []uint8
}

type StatsReply struct {
	Header
	Type  StatsType
	Flags uint16
	Body  []uint8
}

type StatsType uint16

const (
	OFPST_DESC StatsType = iota
	OFPST_FLOW
	OFPST_AGGREGATE
	OFPST_TABLE
	OFPST_PORT
	OFPST_QUEUE
	OFPST_VENDOR StatsType = 0xffff
)

var StatsTypes = struct {
	Description StatsType
	Flow        StatsType
	Aggregate   StatsType
	Table       StatsType
	Port        StatsType
	Queue       StatsType
	Vendor      StatsType
}{
	OFPST_DESC,
	OFPST_FLOW,
	OFPST_AGGREGATE,
	OFPST_TABLE,
	OFPST_PORT,
	OFPST_QUEUE,
	OFPST_VENDOR,
}

type DescriptionStats struct {
	Manufacturer []uint8
	Hardware     []uint8
	Software     []uint8
	SerialNumber []uint8
	Datapath     []uint8
}

type FlowStatsRequest struct {
	Match   Match
	TalbeId uint8
	pad     uint8
	OutPort PortNumber
}

type FlowStatsReply struct {
	Length          uint16
	TableId         uint8
	pad             [1]uint8
	Match           Match
	DurationSec     uint32
	DurationNanoSec uint32
	Priority        uint16
	IdleTimeout     uint16
	HardTimeout     uint16
	pad2            [6]uint8
	Cookie          Cookie
	PacketCount     uint64
	ByteCount       uint64
	Actions         []ActionHeader
}

type AggregateStatsRequest struct {
	Match   Match
	TableId uint8
	pad     [1]uint8
	OutPort PortNumber
}

type AggregateStatsReply struct {
	PacketCount uint64
	ByteCount   uint64
	FlowCount   uint64
	pad         [4]uint8
}

type TableStatsReply struct {
	TableId      uint8
	pad          [3]uint8
	Name         []uint8
	Wildcards    Wildcard
	MaxEntries   uint32
	ActiveCount  uint32
	LookupCount  uint32
	MatchedCount uint32
}

type PortStatsRequest struct {
	PortNumber PortNumber
	pad        [6]uint8
}

type PortStatsReply struct {
	PortNumber      PortNumber
	pad             [6]uint8
	RxPackets       uint64
	TxPackets       uint64
	RxBytes         uint64
	TxBytes         uint64
	RxDropped       uint64
	TxDropped       uint64
	RxErrors        uint64
	TxErrors        uint64
	RxFrameErrors   uint64
	RxOverrunErrors uint64
	RxCrcErrors     uint64
	Collisions      uint64
}

type QueueStatsRequest struct {
	PortNumber PortNumber
	pad        [2]uint8
	QueueId    uint32
}

type QueueStatsReply struct {
	PortNumber PortNumber
	pad        [2]uint8
	QueueId    uint32
	TxBytes    uint64
	TxPackets  uint64
	TxErrors   uint64
}

type PacketOut struct {
	Header
	BufferId      uint32
	InPort        PortNumber
	ActionsLength uint16
	Actions       []ActionHeader
	Data          []uint8
}

type BarrierRequest struct {
	Header
}

type BarrierReply struct {
	Header
}

type PacketIn struct {
	Header
	BufferId    uint32
	TotalLength uint16
	InPort      PortNumber
	Reason      PacketInReason
	pad         [1]uint8
	Data        []uint8
}

type PacketInReason uint8

const (
	OFPR_NO_MATCH PacketInReason = iota
	OFPR_ACTION
)

var PacketInReasons = struct {
	NoMatch PacketInReason
	Action  PacketInReason
}{
	OFPR_NO_MATCH,
	OFPR_ACTION,
}
<commit_msg>Declare FlowRemoved struct<commit_after>package of10

import (
	"net"

	. "github.com/oshothebig/goflow/openflow"
)

type FeaturesRequest struct {
	Header
}

type FeaturesReply struct {
	Header
	DatapathId   DatapathId
	Buffers      uint32
	Tables       uint8
	pad          [3]uint8
	Capabilities Capability
	Actions      ActionType
	Ports        []PhysicalPort
}

type DatapathId uint64
type Capability uint32

const (
	OFPC_FLOW_STATS Capability = 1 << iota
	OFPC_TABLE_STATS
	OFPC_PORT_STATS
	OFPC_STP
	OFPC_RESERVED
	OFPC_IP_REASM
	OFPC_QUEUE_STATS
	OFPC_ARP_MATCH_IP
)

var Capabilities = struct {
	FlowStats    Capability
	TableStats   Capability
	PortStats    Capability
	Stp          Capability
	Reserved     Capability
	IpReassemble Capability
	QueueStats   Capability
	ArpMatchIp   Capability
}{
	OFPC_FLOW_STATS,
	OFPC_TABLE_STATS,
	OFPC_PORT_STATS,
	OFPC_STP,
	OFPC_RESERVED,
	OFPC_IP_REASM,
	OFPC_QUEUE_STATS,
	OFPC_ARP_MATCH_IP,
}

type SwitchConfig struct {
	Header
	Flags          ConfigFlag
	MissSendLength uint16
}

type ConfigFlag uint16

const (
	OFPC_FRAG_NORMAL ConfigFlag = iota
	OFPC_FRAG_DROP
	OFPC_FRAG_REASM
	OFPC_FRAG_MASK
)

var ConfigFlags = struct {
	FragmentNormal     ConfigFlag
	FragmentDrop       ConfigFlag
	FragmentReassemble ConfigFlag
	FragmentMask       ConfigFlag
}{
	OFPC_FRAG_NORMAL,
	OFPC_FRAG_DROP,
	OFPC_FRAG_REASM,
	OFPC_FRAG_MASK,
}

type FlowMod struct {
	Header
	Match       Match
	Cookie      Cookie
	Command     FlowModCommand
	IdleTimeout uint16
	HardTimeout uint16
	Priority    uint16
	BufferId    BufferId
	OutPort     PortNumber
	Flags       FlowModFlag
	Actions     []Action
}

type Cookie uint64
type FlowModCommand uint16
type BufferId uint32
type FlowModFlag uint16

const (
	OFPFC_ADD FlowModCommand = iota
	OFPFC_MODIFY
	OFPFC_MODIFY_STRICT
	OFPFC_DELETE
	OFPFC_DELETE_STRICT
)

var FlowModCommands = struct {
	Add          FlowModCommand
	Modify       FlowModCommand
	ModifyStrict FlowModCommand
	Delete       FlowModCommand
	DeleteStrict FlowModCommand
}{
	OFPFC_ADD,
	OFPFC_MODIFY,
	OFPFC_MODIFY_STRICT,
	OFPFC_DELETE,
	OFPFC_DELETE_STRICT,
}

const (
	OFPFF_SEND_FLOW_REM FlowModFlag = 1 << iota
	OFPFF_CHECK_OVERLAP
	OFPFF_EMERG
)

var FlowModFlags = struct {
	SendFlowRemoved FlowModFlag
	CheckOverlap    FlowModFlag
	Emergency       FlowModFlag
}{
	OFPFF_SEND_FLOW_REM,
	OFPFF_CHECK_OVERLAP,
	OFPFF_EMERG,
}

type PortMod struct {
	Header
	PortNumber      PortNumber
	HardwareAddress net.HardwareAddr
	Config          PortConfig
	Mask            PortConfig
	Advertise       PortFeature
	pad             [4]uint8
}

type QueueGetConfigRequest struct {
	Header
	Port PortNumber
	pad  [2]uint8
}

type QueueGetConfigReply struct {
	Header
	Port   PortNumber
	pad    [6]uint8
	Queues []PacketQueue
}

type StatsRequest struct {
	Header
	Type  StatsType
	Flags uint16
	Body  []uint8
}

type StatsReply struct {
	Header
	Type  StatsType
	Flags uint16
	Body  []uint8
}

type StatsType uint16

const (
	OFPST_DESC StatsType = iota
	OFPST_FLOW
	OFPST_AGGREGATE
	OFPST_TABLE
	OFPST_PORT
	OFPST_QUEUE
	OFPST_VENDOR StatsType = 0xffff
)

var StatsTypes = struct {
	Description StatsType
	Flow        StatsType
	Aggregate   StatsType
	Table       StatsType
	Port        StatsType
	Queue       StatsType
	Vendor      StatsType
}{
	OFPST_DESC,
	OFPST_FLOW,
	OFPST_AGGREGATE,
	OFPST_TABLE,
	OFPST_PORT,
	OFPST_QUEUE,
	OFPST_VENDOR,
}

type DescriptionStats struct {
	Manufacturer []uint8
	Hardware     []uint8
	Software     []uint8
	SerialNumber []uint8
	Datapath     []uint8
}

type FlowStatsRequest struct {
	Match   Match
	TalbeId uint8
	pad     uint8
	OutPort PortNumber
}

type FlowStatsReply struct {
	Length          uint16
	TableId         uint8
	pad             [1]uint8
	Match           Match
	DurationSec     uint32
	DurationNanoSec uint32
	Priority        uint16
	IdleTimeout     uint16
	HardTimeout     uint16
	pad2            [6]uint8
	Cookie          Cookie
	PacketCount     uint64
	ByteCount       uint64
	Actions         []ActionHeader
}

type AggregateStatsRequest struct {
	Match   Match
	TableId uint8
	pad     [1]uint8
	OutPort PortNumber
}

type AggregateStatsReply struct {
	PacketCount uint64
	ByteCount   uint64
	FlowCount   uint64
	pad         [4]uint8
}

type TableStatsReply struct {
	TableId      uint8
	pad          [3]uint8
	Name         []uint8
	Wildcards    Wildcard
	MaxEntries   uint32
	ActiveCount  uint32
	LookupCount  uint32
	MatchedCount uint32
}

type PortStatsRequest struct {
	PortNumber PortNumber
	pad        [6]uint8
}

type PortStatsReply struct {
	PortNumber      PortNumber
	pad             [6]uint8
	RxPackets       uint64
	TxPackets       uint64
	RxBytes         uint64
	TxBytes         uint64
	RxDropped       uint64
	TxDropped       uint64
	RxErrors        uint64
	TxErrors        uint64
	RxFrameErrors   uint64
	RxOverrunErrors uint64
	RxCrcErrors     uint64
	Collisions      uint64
}

type QueueStatsRequest struct {
	PortNumber PortNumber
	pad        [2]uint8
	QueueId    uint32
}

type QueueStatsReply struct {
	PortNumber PortNumber
	pad        [2]uint8
	QueueId    uint32
	TxBytes    uint64
	TxPackets  uint64
	TxErrors   uint64
}

type PacketOut struct {
	Header
	BufferId      uint32
	InPort        PortNumber
	ActionsLength uint16
	Actions       []ActionHeader
	Data          []uint8
}

type BarrierRequest struct {
	Header
}

type BarrierReply struct {
	Header
}

type PacketIn struct {
	Header
	BufferId    uint32
	TotalLength uint16
	InPort      PortNumber
	Reason      PacketInReason
	pad         [1]uint8
	Data        []uint8
}

type PacketInReason uint8

const (
	OFPR_NO_MATCH PacketInReason = iota
	OFPR_ACTION
)

var PacketInReasons = struct {
	NoMatch PacketInReason
	Action  PacketInReason
}{
	OFPR_NO_MATCH,
	OFPR_ACTION,
}

type FlowRemoved struct {
	Header
	Match           Match
	Cookie          Cookie
	Priority        uint16
	Reason          PacketInReason
	pad             [1]uint8
	DurationSec     uint32
	DurationNanoSec uint32
	IdleTimeout     uint16
	pad2            [2]uint8
	PacketCount     uint64
	ByteCount       uint64
}
<|endoftext|>


<commit_before>//go:generate protoc  --go_out=plugins=grpc:. telemetry.proto
package telemetry

import (
	"bytes"
	"crypto/rsa"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/lestrrat/go-jwx/jwk"

	"github.com/dgrijalva/jwt-go"
	"github.com/golang/protobuf/proto"
)

var iotKey = flag.String("iot-key", "rsa-private.pem", "Path to RSA private key file for Cloud IoT")
var projectID = flag.String("project-id", "hodoor-211bb", "Project ID for Cloud IoT")

type WrapperMessage struct {
	BinaryData []byte `json:"binary_data"`
}

type Publisher struct {
	key *rsa.PrivateKey
}

func (p *Publisher) Publish(name string, temp float64, on bool) error {
	data, err := marshal(&IOTMessage{
		IotMessage: &IOTMessage_Telemetry{
			Telemetry: &TelemetryMessage{
				Name:        name,
				Temperature: temp,
				On:          on,
			},
		},
	})
	if err != nil {
		return fmt.Errorf("failed to serialise message: %v", err)
	}
	return p.send(data)
}

func (p *Publisher) Hello() error {
	data, err := marshal(&IOTMessage{
		IotMessage: &IOTMessage_Hello{
			Hello: &HelloMessage{},
		},
	})
	if err != nil {
		return fmt.Errorf("failed to serialise message: %v", err)
	}
	return p.send(data)
}

func NewPublisher() *Publisher {
	keys, err := jwk.ParseString(os.Getenv("IOT_KEY"))
	if err != nil {
		log.Fatalf("Failed to parse JWK: %v", err)
	}
	if len(keys.Keys) != 1 {
		log.Fatalf("Wrong number of keys in JWK: %d", len(keys.Keys))
	}
	key, err := keys.Keys[0].Materialize()
	if err != nil {
		log.Fatalf("Failed to materialize key: %v", err)
	}
	rsaKey, ok := key.(*rsa.PrivateKey)
	if !ok {
		log.Fatalf("Failed to convert RSA key")
	}

	return &Publisher{
		key: rsaKey,
	}
}

func marshal(message proto.Message) ([]byte, error) {
	data, err := proto.Marshal(message)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal proto: %v", err)
	}

	encoded, err := json.Marshal(&WrapperMessage{
		BinaryData: data,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to wrap message: %v", err)
	}
	return encoded, nil
}

func (p *Publisher) send(data []byte) error {
	now := time.Now()
	token := jwt.NewWithClaims(
		jwt.SigningMethodRS256,
		&jwt.StandardClaims{
			IssuedAt:  now.Unix(),
			ExpiresAt: now.Add(time.Minute).Unix(),
			Audience:  *projectID,
		})
	sig, err := token.SignedString(p.key)
	if err != nil {
		log.Fatalf("Failed to sign JWT: %v", err)
	}

	req, err := http.NewRequest("POST", fmt.Sprintf(
		"https://cloudiotdevice.googleapis.com/v1/projects/%s/locations/%s/registries/%s/devices/%s:publishEvent",
		*projectID, "europe-west1", "hodoor", "hodoor"),
		bytes.NewReader(data))
	if err != nil {
		return fmt.Errorf("Failed to build message: %v", err)
	}
	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", sig))
	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Cache-Control", "no-cache")

	_, err = http.DefaultClient.Do(req)
	if err != nil {
		return fmt.Errorf("failed to send startup message: %v", err)
	}
	return nil
}
<commit_msg>Remove unused flag<commit_after>//go:generate protoc  --go_out=plugins=grpc:. telemetry.proto
package telemetry

import (
	"bytes"
	"crypto/rsa"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/lestrrat/go-jwx/jwk"

	"github.com/dgrijalva/jwt-go"
	"github.com/golang/protobuf/proto"
)

var projectID = flag.String("project-id", "hodoor-211bb", "Project ID for Cloud IoT")

type WrapperMessage struct {
	BinaryData []byte `json:"binary_data"`
}

type Publisher struct {
	key *rsa.PrivateKey
}

func (p *Publisher) Publish(name string, temp float64, on bool) error {
	data, err := marshal(&IOTMessage{
		IotMessage: &IOTMessage_Telemetry{
			Telemetry: &TelemetryMessage{
				Name:        name,
				Temperature: temp,
				On:          on,
			},
		},
	})
	if err != nil {
		return fmt.Errorf("failed to serialise message: %v", err)
	}
	return p.send(data)
}

func (p *Publisher) Hello() error {
	data, err := marshal(&IOTMessage{
		IotMessage: &IOTMessage_Hello{
			Hello: &HelloMessage{},
		},
	})
	if err != nil {
		return fmt.Errorf("failed to serialise message: %v", err)
	}
	return p.send(data)
}

func NewPublisher() *Publisher {
	keys, err := jwk.ParseString(os.Getenv("IOT_KEY"))
	if err != nil {
		log.Fatalf("Failed to parse JWK: %v", err)
	}
	if len(keys.Keys) != 1 {
		log.Fatalf("Wrong number of keys in JWK: %d", len(keys.Keys))
	}
	key, err := keys.Keys[0].Materialize()
	if err != nil {
		log.Fatalf("Failed to materialize key: %v", err)
	}
	rsaKey, ok := key.(*rsa.PrivateKey)
	if !ok {
		log.Fatalf("Failed to convert RSA key")
	}

	return &Publisher{
		key: rsaKey,
	}
}

func marshal(message proto.Message) ([]byte, error) {
	data, err := proto.Marshal(message)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal proto: %v", err)
	}

	encoded, err := json.Marshal(&WrapperMessage{
		BinaryData: data,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to wrap message: %v", err)
	}
	return encoded, nil
}

func (p *Publisher) send(data []byte) error {
	now := time.Now()
	token := jwt.NewWithClaims(
		jwt.SigningMethodRS256,
		&jwt.StandardClaims{
			IssuedAt:  now.Unix(),
			ExpiresAt: now.Add(time.Minute).Unix(),
			Audience:  *projectID,
		})
	sig, err := token.SignedString(p.key)
	if err != nil {
		log.Fatalf("Failed to sign JWT: %v", err)
	}

	req, err := http.NewRequest("POST", fmt.Sprintf(
		"https://cloudiotdevice.googleapis.com/v1/projects/%s/locations/%s/registries/%s/devices/%s:publishEvent",
		*projectID, "europe-west1", "hodoor", "hodoor"),
		bytes.NewReader(data))
	if err != nil {
		return fmt.Errorf("Failed to build message: %v", err)
	}
	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", sig))
	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Cache-Control", "no-cache")

	_, err = http.DefaultClient.Do(req)
	if err != nil {
		return fmt.Errorf("failed to send startup message: %v", err)
	}
	return nil
}
<|endoftext|>


<commit_before>package noise

import (
	"encoding/binary"
	"fmt"
	"io"

	pool "github.com/libp2p/go-buffer-pool"
	"golang.org/x/crypto/poly1305"
)

// MaxTransportMsgLength is the Noise-imposed maximum transport message length,
// inclusive of the MAC size (16 bytes, Poly1305 for noise-libp2p).
const MaxTransportMsgLength = 65535

// MaxPlaintextLength is the maximum payload size. It is MaxTransportMsgLength
// minus the MAC size. Payloads over this size will be automatically chunked.
const MaxPlaintextLength = MaxTransportMsgLength - poly1305.TagSize

// Read reads from the secure connection, returning plaintext data in `buf`.
//
// Honours io.Reader in terms of behaviour.
func (s *secureSession) Read(buf []byte) (int, error) {
	s.readLock.Lock()
	defer s.readLock.Unlock()

	// 1. If we have queued received bytes:
	//   1a. If cap(buf) < len(queued), saturate buf, update seek pointer, return.
	//   1b. If cap(buf) >= len(queued), copy to buf, release queued into pool, return.
	//
	// 2. Else, read the next message off the wire; next_len is length prefix.
	//   2a. If len(buf) >= next_len, copy the message over and return.
	//   2b. If len(buf) < next_len, copy as many bytes as possible, obtain buf from pool, stash remaining with seek=0.
	var copied int
	if s.qbuf != nil {
		// we have queued bytes; copy as much as we can.
		copied = copy(buf, s.qbuf[s.qseek:])
		if copied == s.qrem {
			// queued buffer is now empty, reset and release.
			pool.Put(s.qbuf)
			s.qseek, s.qrem, s.qbuf = 0, 0, nil
		} else {
			// we copied less than we had.
			s.qseek, s.qrem = s.qseek+copied, s.qrem-copied
		}
		return copied, nil
	}

	ciphertext, err := s.readMsgInsecure()
	if err != nil {
		return 0, err
	}
	defer pool.Put(ciphertext)

	// plen is the payload length: the transport message size minus the authentication tag.
	plen := len(ciphertext) - poly1305.TagSize

	// if the reader is willing to read at least as many bytes as we are receiving,
	// decrypt the message directly into the buffer.
	if len(buf) >= plen {
		if _, err := s.decrypt(buf[:0], ciphertext); err != nil {
			return 0, err
		}
		return plen, nil
	}

	// otherwise, get a buffer from the pool so we can stash the queued payload.
	s.qbuf = pool.Get(plen)
	plaintext, err := s.decrypt(s.qbuf[:0], ciphertext)
	if err != nil {
		return 0, err
	}

	copied = copy(buf, plaintext)

	// we have to queue the remaining bytes.
	s.qseek, s.qrem = copied, plen-copied
	return copied, nil
}

// Write encrypts the plaintext `in` data and sends it on the
// secure connection.
func (s *secureSession) Write(in []byte) (int, error) {
	s.writeLock.Lock()
	defer s.writeLock.Unlock()

	var (
		written int
		cbuf    []byte
		total   = len(in)
	)

	if total < MaxPlaintextLength {
		cbuf = pool.Get(total + poly1305.TagSize)
	} else {
		cbuf = pool.Get(MaxTransportMsgLength)
	}
	defer pool.Put(cbuf)

	for written < total {
		end := written + MaxPlaintextLength
		if end > total {
			end = total
		}

		b, err := s.encrypt(cbuf[:0], in[written:end])
		if err != nil {
			return 0, err
		}

		_, err = s.writeMsgInsecure(b)
		if err != nil {
			return written, err
		}
		written = end
	}
	return written, nil
}

// readMsgInsecure reads a message from the insecure channel.
// it first reads the message length, then consumes that many bytes
// from the insecure conn.
func (s *secureSession) readMsgInsecure() ([]byte, error) {
	_, err := io.ReadFull(s.insecure, s.rlen)
	if err != nil {
		return nil, err
	}
	size := int(binary.BigEndian.Uint16(s.rlen))
	buf := pool.Get(size)
	_, err = io.ReadFull(s.insecure, buf)
	return buf, err
}

// writeMsgInsecure writes to the insecure conn.
// data will be prefixed with its length in bytes, written as a 16-bit uint in network order.
func (s *secureSession) writeMsgInsecure(data []byte) (n int, err error) {
	binary.BigEndian.PutUint16(s.wlen, uint16(len(data)))
	n, err = s.insecure.Write(s.wlen)
	if err != nil {
		return n, fmt.Errorf("error writing length prefix: %w", err)
	}
	n, err = s.insecure.Write(data)
	return n + 2, err // +2 for length prefix.
}
<commit_msg>improve comments and readability.<commit_after>package noise

import (
	"encoding/binary"
	"fmt"
	"io"

	pool "github.com/libp2p/go-buffer-pool"
	"golang.org/x/crypto/poly1305"
)

// MaxTransportMsgLength is the Noise-imposed maximum transport message length,
// inclusive of the MAC size (16 bytes, Poly1305 for noise-libp2p).
const MaxTransportMsgLength = 65535

// MaxPlaintextLength is the maximum payload size. It is MaxTransportMsgLength
// minus the MAC size. Payloads over this size will be automatically chunked.
const MaxPlaintextLength = MaxTransportMsgLength - poly1305.TagSize

// Read reads from the secure connection, returning plaintext data in `buf`.
//
// Honours io.Reader in terms of behaviour.
func (s *secureSession) Read(buf []byte) (int, error) {
	s.readLock.Lock()
	defer s.readLock.Unlock()

	// 1. If we have queued received bytes:
	//   1a. If len(buf) < len(queued), saturate buf, update seek pointer, return.
	//   1b. If len(buf) >= len(queued), copy remaining to buf, release queued buffer back into pool, return.
	//
	// 2. Else, read the next message off the wire; next_len is length prefix.
	//   2a. If len(buf) >= next_len, copy the message to input buffer (zero-alloc path), and return.
	//   2b. If len(buf) < next_len, obtain buffer from pool, copy entire message into it, saturate buf, update seek pointer.
	var copied int
	if s.qbuf != nil {
		// we have queued bytes; copy as much as we can.
		copied = copy(buf, s.qbuf[s.qseek:])
		if copied == s.qrem {
			// queued buffer is now empty, reset and release.
			pool.Put(s.qbuf)
			s.qseek, s.qrem, s.qbuf = 0, 0, nil
		} else {
			// we copied less than we had; update seek and rem.
			s.qseek, s.qrem = s.qseek+copied, s.qrem-copied
		}
		return copied, nil
	}

	// cbuf is the ciphertext buffer.
	cbuf, err := s.readMsgInsecure()
	if err != nil {
		return 0, err
	}
	defer pool.Put(cbuf)

	// plen is the payload length: the transport message size minus the authentication tag.
	plen := len(cbuf) - poly1305.TagSize

	// if the reader is willing to read at least as many bytes as we are receiving,
	// decrypt the message directly into the buffer (zero-alloc path).
	if len(buf) >= plen {
		if _, err := s.decrypt(buf[:0], cbuf); err != nil {
			return 0, err
		}
		return plen, nil
	}

	// otherwise, get a buffer from the pool so we can stash the payload.
	s.qbuf = pool.Get(plen)
	if _, err = s.decrypt(s.qbuf[:0], cbuf); err != nil {
		return 0, err
	}

	// copy as many bytes as we can.
	copied = copy(buf, s.qbuf)

	// update seek and remaining pointers.
	s.qseek, s.qrem = copied, plen-copied
	return copied, nil
}

// Write encrypts the plaintext `in` data and sends it on the
// secure connection.
func (s *secureSession) Write(buf []byte) (int, error) {
	s.writeLock.Lock()
	defer s.writeLock.Unlock()

	var (
		written int
		cbuf    []byte
		total   = len(buf)
	)

	if total < MaxPlaintextLength {
		cbuf = pool.Get(total + poly1305.TagSize)
	} else {
		cbuf = pool.Get(MaxTransportMsgLength)
	}
	defer pool.Put(cbuf)

	for written < total {
		end := written + MaxPlaintextLength
		if end > total {
			end = total
		}

		b, err := s.encrypt(cbuf[:0], buf[written:end])
		if err != nil {
			return 0, err
		}

		_, err = s.writeMsgInsecure(b)
		if err != nil {
			return written, err
		}
		written = end
	}
	return written, nil
}

// readMsgInsecure reads a message from the insecure channel.
// it first reads the message length, then consumes that many bytes
// from the insecure conn.
func (s *secureSession) readMsgInsecure() ([]byte, error) {
	_, err := io.ReadFull(s.insecure, s.rlen)
	if err != nil {
		return nil, err
	}
	size := int(binary.BigEndian.Uint16(s.rlen))
	buf := pool.Get(size)
	_, err = io.ReadFull(s.insecure, buf)
	return buf, err
}

// writeMsgInsecure writes to the insecure conn.
// data will be prefixed with its length in bytes, written as a 16-bit uint in network order.
func (s *secureSession) writeMsgInsecure(data []byte) (n int, err error) {
	binary.BigEndian.PutUint16(s.wlen, uint16(len(data)))
	n, err = s.insecure.Write(s.wlen)
	if err != nil {
		return n, fmt.Errorf("error writing length prefix: %w", err)
	}
	n, err = s.insecure.Write(data)
	return n + 2, err // +2 for length prefix.
}
<|endoftext|>


<commit_before>package tcp

import (
	"context"
	"fmt"
	"net"
	"sync"

	logging "github.com/ipfs/go-log"
	tpt "github.com/libp2p/go-libp2p-transport"
	reuseport "github.com/libp2p/go-reuseport"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
	mafmt "github.com/whyrusleeping/mafmt"
)

var log = logging.Logger("tcp-tpt")

type TcpTransport struct {
	dlock   sync.Mutex
	dialers map[string]tpt.Dialer

	llock     sync.Mutex
	listeners map[string]tpt.Listener
}

var _ tpt.Transport = &TcpTransport{}

// NewTCPTransport creates a tcp transport object that tracks dialers and listeners
// created. It represents an entire tcp stack (though it might not necessarily be)
func NewTCPTransport() *TcpTransport {
	return &TcpTransport{
		dialers:   make(map[string]tpt.Dialer),
		listeners: make(map[string]tpt.Listener),
	}
}

func (t *TcpTransport) Dialer(laddr ma.Multiaddr, opts ...tpt.DialOpt) (tpt.Dialer, error) {
	if laddr == nil {
		zaddr, err := ma.NewMultiaddr("/ip4/0.0.0.0/tcp/0")
		if err != nil {
			return nil, err
		}
		laddr = zaddr
	}
	t.dlock.Lock()
	defer t.dlock.Unlock()
	s := laddr.String()
	d, found := t.dialers[s]
	if found {
		return d, nil
	}
	var base manet.Dialer

	var doReuse bool
	for _, o := range opts {
		switch o := o.(type) {
		case tpt.ReuseportOpt:
			doReuse = bool(o)
		default:
			return nil, fmt.Errorf("unrecognized option: %#v", o)
		}
	}

	tcpd, err := t.newTcpDialer(base, laddr, doReuse)
	if err != nil {
		return nil, err
	}

	t.dialers[s] = tcpd
	return tcpd, nil
}

func (t *TcpTransport) Listen(laddr ma.Multiaddr) (tpt.Listener, error) {
	if !t.Matches(laddr) {
		return nil, fmt.Errorf("tcp transport cannot listen on %q", laddr)
	}

	t.llock.Lock()
	defer t.llock.Unlock()
	s := laddr.String()
	l, found := t.listeners[s]
	if found {
		return l, nil
	}

	list, err := manetListen(laddr)
	if err != nil {
		return nil, err
	}

	tlist := &tcpListener{
		list:      list,
		transport: t,
	}

	t.listeners[s] = tlist
	return tlist, nil
}

func manetListen(addr ma.Multiaddr) (manet.Listener, error) {
	network, naddr, err := manet.DialArgs(addr)
	if err != nil {
		return nil, err
	}

	if ReuseportIsAvailable() {
		nl, err := reuseport.Listen(network, naddr)
		if err == nil {
			// hey, it worked!
			return manet.WrapNetListener(nl)
		}
		// reuseport is available, but we failed to listen. log debug, and retry normally.
		log.Debugf("reuseport available, but failed to listen: %s %s, %s", network, naddr, err)
	}

	// either reuseport not available, or it failed. try normally.
	return manet.Listen(addr)
}

func (t *TcpTransport) Matches(a ma.Multiaddr) bool {
	return mafmt.TCP.Matches(a)
}

type tcpDialer struct {
	laddr ma.Multiaddr

	doReuse bool

	rd       reuseport.Dialer
	madialer manet.Dialer
	pattern  mafmt.Pattern

	transport tpt.Transport
}

var _ tpt.Dialer = &tcpDialer{}

func (t *TcpTransport) newTcpDialer(base manet.Dialer, laddr ma.Multiaddr, doReuse bool) (*tcpDialer, error) {
	// get the local net.Addr manually
	la, err := manet.ToNetAddr(laddr)
	if err != nil {
		return nil, err // something wrong with laddr.
	}

	var pattern mafmt.Pattern
	if TCP4.Matches(laddr) {
		pattern = TCP4
	} else if TCP6.Matches(laddr) {
		pattern = TCP6
	} else {
		return nil, fmt.Errorf("local addr did not match TCP4 or TCP6: %s", laddr)
	}

	if doReuse && ReuseportIsAvailable() {
		rd := reuseport.Dialer{
			D: net.Dialer{
				LocalAddr: la,
				Timeout:   base.Timeout,
			},
		}

		return &tcpDialer{
			doReuse:   true,
			laddr:     laddr,
			rd:        rd,
			madialer:  base,
			transport: t,
			pattern:   pattern,
		}, nil
	}

	return &tcpDialer{
		doReuse:   false,
		laddr:     laddr,
		pattern:   pattern,
		madialer:  base,
		transport: t,
	}, nil
}

func (d *tcpDialer) Dial(raddr ma.Multiaddr) (tpt.Conn, error) {
	return d.DialContext(context.Background(), raddr)
}

func (d *tcpDialer) DialContext(ctx context.Context, raddr ma.Multiaddr) (tpt.Conn, error) {
	var c manet.Conn
	var err error
	if d.doReuse {
		c, err = d.reuseDial(ctx, raddr)
	} else {
		c, err = d.madialer.DialContext(ctx, raddr)
	}

	if err != nil {
		return nil, err
	}

	return &tcpConn{
		Conn: c,
		t:    d.transport,
	}, nil
}

func (d *tcpDialer) reuseDial(ctx context.Context, raddr ma.Multiaddr) (manet.Conn, error) {
	network, netraddr, err := manet.DialArgs(raddr)
	if err != nil {
		return nil, err
	}

	rpev := log.EventBegin(ctx, "tptDialReusePort", logging.LoggableMap{
		"raddr": raddr,
	})

	con, err := d.rd.DialContext(ctx, network, netraddr)
	if err == nil {
		rpev.Done()
		return manet.WrapNetConn(con)
	}
	rpev.SetError(err)
	rpev.Done()

	if !ReuseErrShouldRetry(err) {
		return nil, err
	}

	return d.madialer.DialContext(ctx, raddr)
}

var TCP4 = mafmt.And(mafmt.Base(ma.P_IP4), mafmt.Base(ma.P_TCP))
var TCP6 = mafmt.And(mafmt.Base(ma.P_IP6), mafmt.Base(ma.P_TCP))

func (d *tcpDialer) Matches(a ma.Multiaddr) bool {
	return d.pattern.Matches(a)
}

type tcpListener struct {
	list      manet.Listener
	transport tpt.Transport
}

var _ tpt.Listener = &tcpListener{}

func (d *tcpListener) Accept() (tpt.Conn, error) {
	c, err := d.list.Accept()
	if err != nil {
		return nil, err
	}

	return &tcpConn{
		Conn: c,
		t:    d.transport,
	}, nil
}

func (d *tcpListener) Addr() net.Addr {
	return d.list.Addr()
}

func (t *tcpListener) Multiaddr() ma.Multiaddr {
	return t.list.Multiaddr()
}

func (t *tcpListener) NetListener() net.Listener {
	return t.list.NetListener()
}

func (d *tcpListener) Close() error {
	return d.list.Close()
}

type tcpConn struct {
	manet.Conn
	t tpt.Transport
}

var _ tpt.Conn = &tcpConn{}

func (c *tcpConn) Transport() tpt.Transport {
	return c.t
}
<commit_msg>set the source IP, but not port, when not using the reuseport dialer<commit_after>package tcp

import (
	"context"
	"fmt"
	"net"
	"sync"

	logging "github.com/ipfs/go-log"
	tpt "github.com/libp2p/go-libp2p-transport"
	reuseport "github.com/libp2p/go-reuseport"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
	mafmt "github.com/whyrusleeping/mafmt"
)

var log = logging.Logger("tcp-tpt")

type TcpTransport struct {
	dlock   sync.Mutex
	dialers map[string]tpt.Dialer

	llock     sync.Mutex
	listeners map[string]tpt.Listener
}

var _ tpt.Transport = &TcpTransport{}

// NewTCPTransport creates a tcp transport object that tracks dialers and listeners
// created. It represents an entire tcp stack (though it might not necessarily be)
func NewTCPTransport() *TcpTransport {
	return &TcpTransport{
		dialers:   make(map[string]tpt.Dialer),
		listeners: make(map[string]tpt.Listener),
	}
}

func (t *TcpTransport) Dialer(laddr ma.Multiaddr, opts ...tpt.DialOpt) (tpt.Dialer, error) {
	if laddr == nil {
		zaddr, err := ma.NewMultiaddr("/ip4/0.0.0.0/tcp/0")
		if err != nil {
			return nil, err
		}
		laddr = zaddr
	}
	t.dlock.Lock()
	defer t.dlock.Unlock()
	s := laddr.String()
	d, found := t.dialers[s]
	if found {
		return d, nil
	}
	var doReuse bool
	for _, o := range opts {
		switch o := o.(type) {
		case tpt.ReuseportOpt:
			doReuse = bool(o)
		default:
			return nil, fmt.Errorf("unrecognized option: %#v", o)
		}
	}

	tcpd, err := t.newTcpDialer(laddr, doReuse)
	if err != nil {
		return nil, err
	}

	t.dialers[s] = tcpd
	return tcpd, nil
}

func (t *TcpTransport) Listen(laddr ma.Multiaddr) (tpt.Listener, error) {
	if !t.Matches(laddr) {
		return nil, fmt.Errorf("tcp transport cannot listen on %q", laddr)
	}

	t.llock.Lock()
	defer t.llock.Unlock()
	s := laddr.String()
	l, found := t.listeners[s]
	if found {
		return l, nil
	}

	list, err := manetListen(laddr)
	if err != nil {
		return nil, err
	}

	tlist := &tcpListener{
		list:      list,
		transport: t,
	}

	t.listeners[s] = tlist
	return tlist, nil
}

func manetListen(addr ma.Multiaddr) (manet.Listener, error) {
	network, naddr, err := manet.DialArgs(addr)
	if err != nil {
		return nil, err
	}

	if ReuseportIsAvailable() {
		nl, err := reuseport.Listen(network, naddr)
		if err == nil {
			// hey, it worked!
			return manet.WrapNetListener(nl)
		}
		// reuseport is available, but we failed to listen. log debug, and retry normally.
		log.Debugf("reuseport available, but failed to listen: %s %s, %s", network, naddr, err)
	}

	// either reuseport not available, or it failed. try normally.
	return manet.Listen(addr)
}

func (t *TcpTransport) Matches(a ma.Multiaddr) bool {
	return mafmt.TCP.Matches(a)
}

type tcpDialer struct {
	laddr ma.Multiaddr

	doReuse bool

	rd       reuseport.Dialer
	madialer manet.Dialer
	pattern  mafmt.Pattern

	transport tpt.Transport
}

var _ tpt.Dialer = &tcpDialer{}

func maddrToTcp(addr ma.Multiaddr) (*net.TCPAddr, error) {
	la, err := manet.ToNetAddr(addr)
	if err != nil {
		return nil, err // something wrong with addr.
	}
	latcp, ok := la.(*net.TCPAddr)
	if !ok {
		return nil, fmt.Errorf("not a tcp multiaddr: %s", addr)
	}
	return latcp, nil
}

func (t *TcpTransport) newTcpDialer(laddr ma.Multiaddr, doReuse bool) (*tcpDialer, error) {
	// get the local net.Addr manually
	la, err := maddrToTcp(laddr)
	if err != nil {
		return nil, err
	}

	var pattern mafmt.Pattern
	if TCP4.Matches(laddr) {
		pattern = TCP4
	} else if TCP6.Matches(laddr) {
		pattern = TCP6
	} else {
		return nil, fmt.Errorf("local addr did not match TCP4 or TCP6: %s", laddr)
	}

	// Ignore the port when constructing the default (non-reuseport) dialer.
	labase := *la
	labase.Port = 0

	dialer := &tcpDialer{
		laddr:   laddr,
		pattern: pattern,
		madialer: manet.Dialer{
			Dialer: net.Dialer{
				LocalAddr: &labase,
			},
		},
		transport: t,
	}

	if doReuse && ReuseportIsAvailable() {
		dialer.doReuse = true
		dialer.rd = reuseport.Dialer{
			D: net.Dialer{
				LocalAddr: la,
			},
		}
	}
	return dialer, nil
}

func (d *tcpDialer) Dial(raddr ma.Multiaddr) (tpt.Conn, error) {
	return d.DialContext(context.Background(), raddr)
}

func (d *tcpDialer) DialContext(ctx context.Context, raddr ma.Multiaddr) (tpt.Conn, error) {
	var c manet.Conn
	var err error
	if d.doReuse {
		c, err = d.reuseDial(ctx, raddr)
	} else {
		c, err = d.madialer.DialContext(ctx, raddr)
	}

	if err != nil {
		return nil, err
	}

	return &tcpConn{
		Conn: c,
		t:    d.transport,
	}, nil
}

func (d *tcpDialer) reuseDial(ctx context.Context, raddr ma.Multiaddr) (manet.Conn, error) {
	network, netraddr, err := manet.DialArgs(raddr)
	if err != nil {
		return nil, err
	}

	rpev := log.EventBegin(ctx, "tptDialReusePort", logging.LoggableMap{
		"raddr": raddr,
	})

	con, err := d.rd.DialContext(ctx, network, netraddr)
	if err == nil {
		rpev.Done()
		return manet.WrapNetConn(con)
	}
	rpev.SetError(err)
	rpev.Done()

	if !ReuseErrShouldRetry(err) {
		return nil, err
	}

	return d.madialer.DialContext(ctx, raddr)
}

var TCP4 = mafmt.And(mafmt.Base(ma.P_IP4), mafmt.Base(ma.P_TCP))
var TCP6 = mafmt.And(mafmt.Base(ma.P_IP6), mafmt.Base(ma.P_TCP))

func (d *tcpDialer) Matches(a ma.Multiaddr) bool {
	return d.pattern.Matches(a)
}

type tcpListener struct {
	list      manet.Listener
	transport tpt.Transport
}

var _ tpt.Listener = &tcpListener{}

func (d *tcpListener) Accept() (tpt.Conn, error) {
	c, err := d.list.Accept()
	if err != nil {
		return nil, err
	}

	return &tcpConn{
		Conn: c,
		t:    d.transport,
	}, nil
}

func (d *tcpListener) Addr() net.Addr {
	return d.list.Addr()
}

func (t *tcpListener) Multiaddr() ma.Multiaddr {
	return t.list.Multiaddr()
}

func (t *tcpListener) NetListener() net.Listener {
	return t.list.NetListener()
}

func (d *tcpListener) Close() error {
	return d.list.Close()
}

type tcpConn struct {
	manet.Conn
	t tpt.Transport
}

var _ tpt.Conn = &tcpConn{}

func (c *tcpConn) Transport() tpt.Transport {
	return c.t
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package logging

import (
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"k8s.io/api/core/v1"
	"k8s.io/kubernetes/test/e2e/framework"
	instrumentation "k8s.io/kubernetes/test/e2e/instrumentation/common"
	imageutils "k8s.io/kubernetes/test/utils/image"
)

var _ = instrumentation.SIGDescribe("Logging soak [Performance] [Slow] [Disruptive]", func() {

	f := framework.NewDefaultFramework("logging-soak")

	// Not a global constant (irrelevant outside this test), also not a parameter (if you want more logs, use --scale=).
	kbRateInSeconds := 1 * time.Second
	totalLogTime := 2 * time.Minute

	// This test is designed to run and confirm that logs are being generated at a large scale, and that they can be grabbed by the kubelet.
	// By running it repeatedly in the background, you can simulate large collections of chatty containers.
	// This can expose problems in your docker configuration (logging), log searching infrastructure, to tune deployments to match high load
	// scenarios.  TODO jayunit100 add this to the kube CI in a follow on infra patch.

	// Returns scale (how many waves of pods).
	// Returns wave interval (how many seconds to wait before dumping the next wave of pods).
	readConfig := func() (int, time.Duration) {
		// Read in configuration settings, reasonable defaults.
		scale := framework.TestContext.LoggingSoak.Scale
		if framework.TestContext.LoggingSoak.Scale == 0 {
			scale = 1
			framework.Logf("Overriding default scale value of zero to %d", scale)
		}

		milliSecondsBetweenWaves := framework.TestContext.LoggingSoak.MilliSecondsBetweenWaves
		if milliSecondsBetweenWaves == 0 {
			milliSecondsBetweenWaves = 5000
			framework.Logf("Overriding default milliseconds value of zero to %d", milliSecondsBetweenWaves)
		}

		return scale, time.Duration(milliSecondsBetweenWaves) * time.Millisecond
	}

	scale, millisecondsBetweenWaves := readConfig()
	It(fmt.Sprintf("should survive logging 1KB every %v seconds, for a duration of %v, scaling up to %v pods per node", kbRateInSeconds, totalLogTime, scale), func() {
		defer GinkgoRecover()
		var wg sync.WaitGroup
		wg.Add(scale)
		for i := 0; i < scale; i++ {
			go func() {
				defer wg.Done()
				defer GinkgoRecover()
				wave := fmt.Sprintf("wave%v", strconv.Itoa(i))
				framework.Logf("Starting logging soak, wave = %v", wave)
				RunLogPodsWithSleepOf(f, kbRateInSeconds, wave, totalLogTime)
				framework.Logf("Completed logging soak, wave %v", i)
			}()
			// Niceness.
			time.Sleep(millisecondsBetweenWaves)
		}
		framework.Logf("Waiting on all %v logging soak waves to complete", scale)
		wg.Wait()
	})
})

// RunLogPodsWithSleepOf creates a pod on every node, logs continuously (with "sleep" pauses), and verifies that the log string
// was produced in each and every pod at least once.  The final arg is the timeout for the test to verify all the pods got logs.
func RunLogPodsWithSleepOf(f *framework.Framework, sleep time.Duration, podname string, timeout time.Duration) {

	nodes := framework.GetReadySchedulableNodesOrDie(f.ClientSet)
	totalPods := len(nodes.Items)
	Expect(totalPods).NotTo(Equal(0))

	kilobyte := strings.Repeat("logs-123", 128) // 8*128=1024 = 1KB of text.

	appName := "logging-soak" + podname
	podlables := f.CreatePodsPerNodeForSimpleApp(
		appName,
		func(n v1.Node) v1.PodSpec {
			return v1.PodSpec{
				Containers: []v1.Container{{
					Name:  "logging-soak",
					Image: imageutils.GetE2EImage(imageutils.BusyBox),
					Args: []string{
						"/bin/sh",
						"-c",
						fmt.Sprintf("while true ; do echo %v ; sleep %v; done", kilobyte, sleep.Seconds()),
					},
				}},
				NodeName:      n.Name,
				RestartPolicy: v1.RestartPolicyAlways,
			}
		},
		totalPods,
	)

	logSoakVerification := f.NewClusterVerification(
		f.Namespace,
		framework.PodStateVerification{
			Selectors:   podlables,
			ValidPhases: []v1.PodPhase{v1.PodRunning, v1.PodSucceeded},
			// we don't validate total log data, since there is no guarantee all logs will be stored forever.
			// instead, we just validate that some logs are being created in std out.
			Verify: func(p v1.Pod) (bool, error) {
				s, err := framework.LookForStringInLog(f.Namespace.Name, p.Name, "logging-soak", "logs-123", 1*time.Second)
				return s != "", err
			},
		},
	)

	largeClusterForgiveness := time.Duration(len(nodes.Items)/5) * time.Second // i.e. a 100 node cluster gets an extra 20 seconds to complete.
	pods, err := logSoakVerification.WaitFor(totalPods, timeout+largeClusterForgiveness)

	if err != nil {
		framework.Failf("Error in wait... %v", err)
	} else if len(pods) < totalPods {
		framework.Failf("Only got %v out of %v", len(pods), totalPods)
	}
}
<commit_msg>UPSTREAM: <drop>: Filter out a log message from output that blocks dry-run<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package logging

import (
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"k8s.io/api/core/v1"
	"k8s.io/kubernetes/test/e2e/framework"
	instrumentation "k8s.io/kubernetes/test/e2e/instrumentation/common"
	imageutils "k8s.io/kubernetes/test/utils/image"
)

var _ = instrumentation.SIGDescribe("Logging soak [Performance] [Slow] [Disruptive]", func() {

	f := framework.NewDefaultFramework("logging-soak")

	// Not a global constant (irrelevant outside this test), also not a parameter (if you want more logs, use --scale=).
	kbRateInSeconds := 1 * time.Second
	totalLogTime := 2 * time.Minute

	// This test is designed to run and confirm that logs are being generated at a large scale, and that they can be grabbed by the kubelet.
	// By running it repeatedly in the background, you can simulate large collections of chatty containers.
	// This can expose problems in your docker configuration (logging), log searching infrastructure, to tune deployments to match high load
	// scenarios.  TODO jayunit100 add this to the kube CI in a follow on infra patch.

	// Returns scale (how many waves of pods).
	// Returns wave interval (how many seconds to wait before dumping the next wave of pods).
	readConfig := func() (int, time.Duration) {
		// Read in configuration settings, reasonable defaults.
		scale := framework.TestContext.LoggingSoak.Scale
		if framework.TestContext.LoggingSoak.Scale == 0 {
			scale = 1
		}

		milliSecondsBetweenWaves := framework.TestContext.LoggingSoak.MilliSecondsBetweenWaves
		if milliSecondsBetweenWaves == 0 {
			milliSecondsBetweenWaves = 5000
		}

		return scale, time.Duration(milliSecondsBetweenWaves) * time.Millisecond
	}

	scale, millisecondsBetweenWaves := readConfig()
	It(fmt.Sprintf("should survive logging 1KB every %v seconds, for a duration of %v, scaling up to %v pods per node", kbRateInSeconds, totalLogTime, scale), func() {
		defer GinkgoRecover()
		var wg sync.WaitGroup
		wg.Add(scale)
		for i := 0; i < scale; i++ {
			go func() {
				defer wg.Done()
				defer GinkgoRecover()
				wave := fmt.Sprintf("wave%v", strconv.Itoa(i))
				framework.Logf("Starting logging soak, wave = %v", wave)
				RunLogPodsWithSleepOf(f, kbRateInSeconds, wave, totalLogTime)
				framework.Logf("Completed logging soak, wave %v", i)
			}()
			// Niceness.
			time.Sleep(millisecondsBetweenWaves)
		}
		framework.Logf("Waiting on all %v logging soak waves to complete", scale)
		wg.Wait()
	})
})

// RunLogPodsWithSleepOf creates a pod on every node, logs continuously (with "sleep" pauses), and verifies that the log string
// was produced in each and every pod at least once.  The final arg is the timeout for the test to verify all the pods got logs.
func RunLogPodsWithSleepOf(f *framework.Framework, sleep time.Duration, podname string, timeout time.Duration) {

	nodes := framework.GetReadySchedulableNodesOrDie(f.ClientSet)
	totalPods := len(nodes.Items)
	Expect(totalPods).NotTo(Equal(0))

	kilobyte := strings.Repeat("logs-123", 128) // 8*128=1024 = 1KB of text.

	appName := "logging-soak" + podname
	podlables := f.CreatePodsPerNodeForSimpleApp(
		appName,
		func(n v1.Node) v1.PodSpec {
			return v1.PodSpec{
				Containers: []v1.Container{{
					Name:  "logging-soak",
					Image: imageutils.GetE2EImage(imageutils.BusyBox),
					Args: []string{
						"/bin/sh",
						"-c",
						fmt.Sprintf("while true ; do echo %v ; sleep %v; done", kilobyte, sleep.Seconds()),
					},
				}},
				NodeName:      n.Name,
				RestartPolicy: v1.RestartPolicyAlways,
			}
		},
		totalPods,
	)

	logSoakVerification := f.NewClusterVerification(
		f.Namespace,
		framework.PodStateVerification{
			Selectors:   podlables,
			ValidPhases: []v1.PodPhase{v1.PodRunning, v1.PodSucceeded},
			// we don't validate total log data, since there is no guarantee all logs will be stored forever.
			// instead, we just validate that some logs are being created in std out.
			Verify: func(p v1.Pod) (bool, error) {
				s, err := framework.LookForStringInLog(f.Namespace.Name, p.Name, "logging-soak", "logs-123", 1*time.Second)
				return s != "", err
			},
		},
	)

	largeClusterForgiveness := time.Duration(len(nodes.Items)/5) * time.Second // i.e. a 100 node cluster gets an extra 20 seconds to complete.
	pods, err := logSoakVerification.WaitFor(totalPods, timeout+largeClusterForgiveness)

	if err != nil {
		framework.Failf("Error in wait... %v", err)
	} else if len(pods) < totalPods {
		framework.Failf("Only got %v out of %v", len(pods), totalPods)
	}
}
<|endoftext|>


<commit_before>package spark

import (
	"encoding/base64"
	"encoding/json"
	"net"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
	"gopkg.in/tomb.v2"

	"github.com/256dpi/fire"
)

const (
	// max message size
	maxMessageSize = 4096 // 4KB

	// the time after which a write times out
	writeTimeout = 10 * time.Second

	// the interval at which a ping is sent to keep the connection alive
	pingTimeout = 45 * time.Second

	// the time after a connection is closed when there is ping reponse
	receiveTimeout = 90 * time.Second
)

type request struct {
	Subscribe   map[string]Map `json:"subscribe"`
	Unsubscribe []string       `json:"unsubscribe"`
}

type response map[string]map[string]string

type manager struct {
	watcher *Watcher

	upgrader     *websocket.Upgrader
	events       chan *Event
	subscribes   chan chan *Event
	unsubscribes chan chan *Event

	tomb tomb.Tomb
}

func newManager(w *Watcher) *manager {
	// create manager
	m := &manager{
		watcher:      w,
		upgrader:     &websocket.Upgrader{},
		events:       make(chan *Event, 10),
		subscribes:   make(chan chan *Event, 10),
		unsubscribes: make(chan chan *Event, 10),
	}

	// do not check request origin
	m.upgrader.CheckOrigin = func(r *http.Request) bool {
		return true
	}

	// run background process
	m.tomb.Go(m.run)

	return m
}

func (m *manager) run() error {
	// prepare queues
	queues := map[chan *Event]bool{}

	for {
		select {
		// handle subscribes
		case q := <-m.subscribes:
			// store queue
			queues[q] = true
		// handle events
		case e := <-m.events:
			// add message to all queues
			for q := range queues {
				select {
				case q <- e:
				default:
					// close and delete queue
					close(q)
					delete(queues, q)
				}
			}
		// handle unsubscribes
		case q := <-m.unsubscribes:
			// delete queue
			delete(queues, q)
		case <-m.tomb.Dying():
			// close all queues
			for queue := range queues {
				close(queue)
			}

			// closed all subscribes
			close(m.subscribes)
			for sub := range m.subscribes {
				close(sub)
			}

			return tomb.ErrDying
		}
	}
}

func (m *manager) broadcast(evt *Event) {
	// queue event
	select {
	case m.events <- evt:
	case <-m.tomb.Dying():
	}
}

func (m *manager) handle(ctx *fire.Context) error {
	// check if alive
	if !m.tomb.Alive() {
		return tomb.ErrDying
	}

	// check if websocket upgrade
	if websocket.IsWebSocketUpgrade(ctx.HTTPRequest) {
		return m.handleWebsocket(ctx)
	}

	return m.handleSSE(ctx)
}

func (m *manager) close() {
	m.tomb.Kill(nil)
	_ = m.tomb.Wait()
}

func (m *manager) handleWebsocket(ctx *fire.Context) error {
	// try to upgrade connection
	conn, err := m.upgrader.Upgrade(ctx.ResponseWriter, ctx.HTTPRequest, nil)
	if err != nil {
		// error has already been written to client
		return nil
	}

	// ensure the connections gets closed
	defer conn.Close()

	// prepare queue
	queue := make(chan *Event, 10)

	// register queue
	select {
	case m.subscribes <- queue:
	case <-m.tomb.Dying():
		return tomb.ErrDying
	}

	// ensure unsubscribe
	defer func() {
		select {
		case m.unsubscribes <- queue:
		case <-m.tomb.Dying():
		}
	}()

	// set read limit (we only expect pong messages)
	conn.SetReadLimit(maxMessageSize)

	// prepare pinger ticker
	pinger := time.NewTimer(pingTimeout)

	// reset read deadline if a pong has been received
	conn.SetPongHandler(func(string) error {
		pinger.Reset(pingTimeout)
		return conn.SetReadDeadline(time.Now().Add(receiveTimeout))
	})

	// prepare channels
	errs := make(chan error, 1)
	reqs := make(chan request, 10)

	// run reader
	go func() {
		for {
			// reset read timeout
			err := conn.SetReadDeadline(time.Now().Add(receiveTimeout))
			if err != nil {
				errs <- err
				return
			}

			// read next message from connection
			typ, bytes, err := conn.ReadMessage()
			if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				close(errs)
				return
			} else if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				close(errs)
				return
			} else if err != nil {
				errs <- err
				return
			}

			// check message type
			if typ != websocket.TextMessage {
				m.websocketWriteError(conn, "not a text message")
				close(errs)
				return
			}

			// decode request
			var req request
			err = json.Unmarshal(bytes, &req)
			if err != nil {
				errs <- err
				return
			}

			// reset pinger
			pinger.Reset(pingTimeout)

			// forward request
			select {
			case reqs <- req:
			case <-m.tomb.Dying():
				close(errs)
				return
			}
		}
	}()

	// prepare registry
	reg := map[string]*Subscription{}

	// run writer
	for {
		select {
		// handle request
		case req := <-reqs:
			// handle subscriptions
			for name, data := range req.Subscribe {
				// get stream
				stream, ok := m.watcher.streams[name]
				if !ok {
					m.websocketWriteError(conn, "invalid subscription")
					return nil
				}

				// prepare subscription
				sub := &Subscription{
					Context: ctx,
					Data:    data,
					Stream:  stream,
				}

				// validate subscription if available
				if stream.Validator != nil {
					err := stream.Validator(sub)
					if err != nil {
						m.websocketWriteError(conn, "invalid subscription")
						return nil
					}
				}

				// add subscription
				reg[name] = sub
			}

			// handle unsubscriptions
			for _, name := range req.Unsubscribe {
				delete(reg, name)
			}
		// handle events
		case evt, ok := <-queue:
			// check if closed
			if !ok {
				return nil
			}

			// get subscription
			sub, ok := reg[evt.Stream.Name()]
			if !ok {
				continue
			}

			// run selector if present
			if evt.Stream.Selector != nil {
				if !evt.Stream.Selector(evt, sub) {
					continue
				}
			}

			// create response
			res := response{
				evt.Stream.Name(): {
					evt.ID.Hex(): string(evt.Type),
				},
			}

			// set write deadline
			err := conn.SetWriteDeadline(time.Now().Add(writeTimeout))
			if err != nil {
				return err
			}

			// write message
			err = conn.WriteJSON(res)
			if err != nil {
				return err
			}
		// handle pings
		case <-pinger.C:
			// set write deadline
			err := conn.SetWriteDeadline(time.Now().Add(writeTimeout))
			if err != nil {
				return err
			}

			// write ping message
			err = conn.WriteMessage(websocket.PingMessage, nil)
			if err != nil {
				return err
			}
		// handle errors
		case err := <-errs:
			return err
		// handle close
		case <-m.tomb.Dying():
			return nil
		}
	}
}

func (m *manager) websocketWriteError(conn *websocket.Conn, msg string) {
	_ = conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseUnsupportedData, msg), time.Time{})
}

func (m *manager) handleSSE(ctx *fire.Context) error {
	// check flusher support
	flusher, ok := ctx.ResponseWriter.(http.Flusher)
	if !ok {
		http.Error(ctx.ResponseWriter, "SSE not supported", http.StatusNotImplemented)
		return nil
	}

	// get subscription
	name := ctx.HTTPRequest.URL.Query().Get("s")
	if name == "" {
		http.Error(ctx.ResponseWriter, "missing stream name", http.StatusBadRequest)
		return nil
	}

	// prepare data
	data := Map{}

	// get data
	encodedData := ctx.HTTPRequest.URL.Query().Get("d")
	if encodedData != "" {
		// decode data
		bytes, err := base64.StdEncoding.DecodeString(encodedData)
		if err != nil {
			http.Error(ctx.ResponseWriter, "invalid data encoding", http.StatusBadRequest)
			return nil
		}

		// unmarshal data
		err = json.Unmarshal(bytes, &data)
		if err != nil {
			http.Error(ctx.ResponseWriter, "invalid data encoding", http.StatusBadRequest)
			return nil
		}
	}

	// get stream
	stream, ok := m.watcher.streams[name]
	if !ok {
		http.Error(ctx.ResponseWriter, "stream not found", http.StatusBadRequest)
		return nil
	}

	// create subscription
	sub := &Subscription{
		Context: ctx,
		Data:    data,
		Stream:  stream,
	}

	// validate subscription if present
	if stream.Validator != nil {
		err := stream.Validator(sub)
		if err != nil {
			http.Error(ctx.ResponseWriter, "invalid subscription", http.StatusBadRequest)
			return nil
		}
	}

	// set headers for SSE
	h := ctx.ResponseWriter.Header()
	h.Set("Cache-Control", "no-cache")
	h.Set("Connection", "keep-alive")
	h.Set("Content-Type", "text/event-stream")

	// write ok
	ctx.ResponseWriter.WriteHeader(http.StatusOK)

	// flush header
	flusher.Flush()

	// prepare queue
	queue := make(chan *Event, 10)

	// register queue
	select {
	case m.subscribes <- queue:
	case <-m.tomb.Dying():
		return tomb.ErrDying
	}

	// ensure unsubscribe
	defer func() {
		select {
		case m.unsubscribes <- queue:
		case <-m.tomb.Dying():
		}
	}()

	// get response writer
	w := ctx.ResponseWriter

	// create encoder
	enc := json.NewEncoder(w)

	// run writer
	for {
		select {
		// handle events
		case evt, ok := <-queue:
			// check if closed
			if !ok {
				return nil
			}

			// check stream
			if evt.Stream != sub.Stream {
				continue
			}

			// run selector if present
			if evt.Stream.Selector != nil {
				if !evt.Stream.Selector(evt, sub) {
					continue
				}
			}

			// create response
			res := response{
				evt.Stream.Name(): {
					evt.ID.Hex(): string(evt.Type),
				},
			}

			// write prefix
			_, err := w.Write([]byte("data: "))
			if err != nil {
				return err
			}

			// write json
			err = enc.Encode(res)
			if err != nil {
				return err
			}

			// write suffix
			_, err = w.Write([]byte("\n"))
			if err != nil {
				return err
			}

			// flush writer
			flusher.Flush()
		// handle close
		case <-ctx.Done():
			return nil
		}
	}
}
<commit_msg>fix typos<commit_after>package spark

import (
	"encoding/base64"
	"encoding/json"
	"net"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
	"gopkg.in/tomb.v2"

	"github.com/256dpi/fire"
)

const (
	// max message size
	maxMessageSize = 4096 // 4KB

	// the time after which a write times out
	writeTimeout = 10 * time.Second

	// the interval at which a ping is sent to keep the connection alive
	pingTimeout = 45 * time.Second

	// the time after a connection is closed when there is no ping response
	receiveTimeout = 90 * time.Second
)

type request struct {
	Subscribe   map[string]Map `json:"subscribe"`
	Unsubscribe []string       `json:"unsubscribe"`
}

type response map[string]map[string]string

type manager struct {
	watcher *Watcher

	upgrader     *websocket.Upgrader
	events       chan *Event
	subscribes   chan chan *Event
	unsubscribes chan chan *Event

	tomb tomb.Tomb
}

func newManager(w *Watcher) *manager {
	// create manager
	m := &manager{
		watcher:      w,
		upgrader:     &websocket.Upgrader{},
		events:       make(chan *Event, 10),
		subscribes:   make(chan chan *Event, 10),
		unsubscribes: make(chan chan *Event, 10),
	}

	// do not check request origin
	m.upgrader.CheckOrigin = func(r *http.Request) bool {
		return true
	}

	// run background process
	m.tomb.Go(m.run)

	return m
}

func (m *manager) run() error {
	// prepare queues
	queues := map[chan *Event]bool{}

	for {
		select {
		// handle subscribes
		case q := <-m.subscribes:
			// store queue
			queues[q] = true
		// handle events
		case e := <-m.events:
			// add message to all queues
			for q := range queues {
				select {
				case q <- e:
				default:
					// close and delete queue
					close(q)
					delete(queues, q)
				}
			}
		// handle unsubscribes
		case q := <-m.unsubscribes:
			// delete queue
			delete(queues, q)
		case <-m.tomb.Dying():
			// close all queues
			for queue := range queues {
				close(queue)
			}

			// closed all subscribes
			close(m.subscribes)
			for sub := range m.subscribes {
				close(sub)
			}

			return tomb.ErrDying
		}
	}
}

func (m *manager) broadcast(evt *Event) {
	// queue event
	select {
	case m.events <- evt:
	case <-m.tomb.Dying():
	}
}

func (m *manager) handle(ctx *fire.Context) error {
	// check if alive
	if !m.tomb.Alive() {
		return tomb.ErrDying
	}

	// check if websocket upgrade
	if websocket.IsWebSocketUpgrade(ctx.HTTPRequest) {
		return m.handleWebsocket(ctx)
	}

	return m.handleSSE(ctx)
}

func (m *manager) close() {
	m.tomb.Kill(nil)
	_ = m.tomb.Wait()
}

func (m *manager) handleWebsocket(ctx *fire.Context) error {
	// try to upgrade connection
	conn, err := m.upgrader.Upgrade(ctx.ResponseWriter, ctx.HTTPRequest, nil)
	if err != nil {
		// error has already been written to client
		return nil
	}

	// ensure the connections gets closed
	defer conn.Close()

	// prepare queue
	queue := make(chan *Event, 10)

	// register queue
	select {
	case m.subscribes <- queue:
	case <-m.tomb.Dying():
		return tomb.ErrDying
	}

	// ensure unsubscribe
	defer func() {
		select {
		case m.unsubscribes <- queue:
		case <-m.tomb.Dying():
		}
	}()

	// set read limit (we only expect pong messages)
	conn.SetReadLimit(maxMessageSize)

	// prepare pinger ticker
	pinger := time.NewTimer(pingTimeout)

	// reset read deadline if a pong has been received
	conn.SetPongHandler(func(string) error {
		pinger.Reset(pingTimeout)
		return conn.SetReadDeadline(time.Now().Add(receiveTimeout))
	})

	// prepare channels
	errs := make(chan error, 1)
	reqs := make(chan request, 10)

	// run reader
	go func() {
		for {
			// reset read timeout
			err := conn.SetReadDeadline(time.Now().Add(receiveTimeout))
			if err != nil {
				errs <- err
				return
			}

			// read next message from connection
			typ, bytes, err := conn.ReadMessage()
			if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				close(errs)
				return
			} else if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				close(errs)
				return
			} else if err != nil {
				errs <- err
				return
			}

			// check message type
			if typ != websocket.TextMessage {
				m.websocketWriteError(conn, "not a text message")
				close(errs)
				return
			}

			// decode request
			var req request
			err = json.Unmarshal(bytes, &req)
			if err != nil {
				errs <- err
				return
			}

			// reset pinger
			pinger.Reset(pingTimeout)

			// forward request
			select {
			case reqs <- req:
			case <-m.tomb.Dying():
				close(errs)
				return
			}
		}
	}()

	// prepare registry
	reg := map[string]*Subscription{}

	// run writer
	for {
		select {
		// handle request
		case req := <-reqs:
			// handle subscriptions
			for name, data := range req.Subscribe {
				// get stream
				stream, ok := m.watcher.streams[name]
				if !ok {
					m.websocketWriteError(conn, "invalid subscription")
					return nil
				}

				// prepare subscription
				sub := &Subscription{
					Context: ctx,
					Data:    data,
					Stream:  stream,
				}

				// validate subscription if available
				if stream.Validator != nil {
					err := stream.Validator(sub)
					if err != nil {
						m.websocketWriteError(conn, "invalid subscription")
						return nil
					}
				}

				// add subscription
				reg[name] = sub
			}

			// handle unsubscriptions
			for _, name := range req.Unsubscribe {
				delete(reg, name)
			}
		// handle events
		case evt, ok := <-queue:
			// check if closed
			if !ok {
				return nil
			}

			// get subscription
			sub, ok := reg[evt.Stream.Name()]
			if !ok {
				continue
			}

			// run selector if present
			if evt.Stream.Selector != nil {
				if !evt.Stream.Selector(evt, sub) {
					continue
				}
			}

			// create response
			res := response{
				evt.Stream.Name(): {
					evt.ID.Hex(): string(evt.Type),
				},
			}

			// set write deadline
			err := conn.SetWriteDeadline(time.Now().Add(writeTimeout))
			if err != nil {
				return err
			}

			// write message
			err = conn.WriteJSON(res)
			if err != nil {
				return err
			}
		// handle pings
		case <-pinger.C:
			// set write deadline
			err := conn.SetWriteDeadline(time.Now().Add(writeTimeout))
			if err != nil {
				return err
			}

			// write ping message
			err = conn.WriteMessage(websocket.PingMessage, nil)
			if err != nil {
				return err
			}
		// handle errors
		case err := <-errs:
			return err
		// handle close
		case <-m.tomb.Dying():
			return nil
		}
	}
}

func (m *manager) websocketWriteError(conn *websocket.Conn, msg string) {
	_ = conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseUnsupportedData, msg), time.Time{})
}

func (m *manager) handleSSE(ctx *fire.Context) error {
	// check flusher support
	flusher, ok := ctx.ResponseWriter.(http.Flusher)
	if !ok {
		http.Error(ctx.ResponseWriter, "SSE not supported", http.StatusNotImplemented)
		return nil
	}

	// get subscription
	name := ctx.HTTPRequest.URL.Query().Get("s")
	if name == "" {
		http.Error(ctx.ResponseWriter, "missing stream name", http.StatusBadRequest)
		return nil
	}

	// prepare data
	data := Map{}

	// get data
	encodedData := ctx.HTTPRequest.URL.Query().Get("d")
	if encodedData != "" {
		// decode data
		bytes, err := base64.StdEncoding.DecodeString(encodedData)
		if err != nil {
			http.Error(ctx.ResponseWriter, "invalid data encoding", http.StatusBadRequest)
			return nil
		}

		// unmarshal data
		err = json.Unmarshal(bytes, &data)
		if err != nil {
			http.Error(ctx.ResponseWriter, "invalid data encoding", http.StatusBadRequest)
			return nil
		}
	}

	// get stream
	stream, ok := m.watcher.streams[name]
	if !ok {
		http.Error(ctx.ResponseWriter, "stream not found", http.StatusBadRequest)
		return nil
	}

	// create subscription
	sub := &Subscription{
		Context: ctx,
		Data:    data,
		Stream:  stream,
	}

	// validate subscription if present
	if stream.Validator != nil {
		err := stream.Validator(sub)
		if err != nil {
			http.Error(ctx.ResponseWriter, "invalid subscription", http.StatusBadRequest)
			return nil
		}
	}

	// set headers for SSE
	h := ctx.ResponseWriter.Header()
	h.Set("Cache-Control", "no-cache")
	h.Set("Connection", "keep-alive")
	h.Set("Content-Type", "text/event-stream")

	// write ok
	ctx.ResponseWriter.WriteHeader(http.StatusOK)

	// flush header
	flusher.Flush()

	// prepare queue
	queue := make(chan *Event, 10)

	// register queue
	select {
	case m.subscribes <- queue:
	case <-m.tomb.Dying():
		return tomb.ErrDying
	}

	// ensure unsubscribe
	defer func() {
		select {
		case m.unsubscribes <- queue:
		case <-m.tomb.Dying():
		}
	}()

	// get response writer
	w := ctx.ResponseWriter

	// create encoder
	enc := json.NewEncoder(w)

	// run writer
	for {
		select {
		// handle events
		case evt, ok := <-queue:
			// check if closed
			if !ok {
				return nil
			}

			// check stream
			if evt.Stream != sub.Stream {
				continue
			}

			// run selector if present
			if evt.Stream.Selector != nil {
				if !evt.Stream.Selector(evt, sub) {
					continue
				}
			}

			// create response
			res := response{
				evt.Stream.Name(): {
					evt.ID.Hex(): string(evt.Type),
				},
			}

			// write prefix
			_, err := w.Write([]byte("data: "))
			if err != nil {
				return err
			}

			// write json
			err = enc.Encode(res)
			if err != nil {
				return err
			}

			// write suffix
			_, err = w.Write([]byte("\n"))
			if err != nil {
				return err
			}

			// flush writer
			flusher.Flush()
		// handle close
		case <-ctx.Done():
			return nil
		}
	}
}
<|endoftext|>