<commit_before>// Package minify relates MIME type to minifiers. Several minifiers are provided in the subpackages.
package minify

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path"
	"regexp"
	"strings"
	"sync"

	"github.com/tdewolff/parse/v2"
	"github.com/tdewolff/parse/v2/buffer"
)

// Warning is used to report usage warnings such as using a deprecated feature
var Warning = log.New(os.Stderr, "WARNING: ", 0)

// ErrNotExist is returned when no minifier exists for a given mimetype.
var ErrNotExist = errors.New("minifier does not exist for mimetype")

// ErrClosedWriter is returned when writing to a closed writer.
var ErrClosedWriter = errors.New("write on closed writer")

////////////////////////////////////////////////////////////////

// MinifierFunc is a function that implements Minifer.
type MinifierFunc func(*M, io.Writer, io.Reader, map[string]string) error

// Minify calls f(m, w, r, params)
func (f MinifierFunc) Minify(m *M, w io.Writer, r io.Reader, params map[string]string) error {
	return f(m, w, r, params)
}

// Minifier is the interface for minifiers.
// The *M parameter is used for minifying embedded resources, such as JS within HTML.
type Minifier interface {
	Minify(*M, io.Writer, io.Reader, map[string]string) error
}

////////////////////////////////////////////////////////////////

type patternMinifier struct {
	pattern *regexp.Regexp
	Minifier
}

type cmdMinifier struct {
	cmd *exec.Cmd
}

var cmdArgExtension = regexp.MustCompile(`^\.[0-9a-zA-Z]+`)

func (c *cmdMinifier) Minify(_ *M, w io.Writer, r io.Reader, _ map[string]string) error {
	cmd := &exec.Cmd{}
	*cmd = *c.cmd // concurrency safety

	var in, out *os.File
	for i, arg := range cmd.Args {
		if j := strings.Index(arg, "$in"); j != -1 {
			var err error
			ext := cmdArgExtension.FindString(arg[j+3:])
			if in, err = ioutil.TempFile("", "minify-in-*"+ext); err != nil {
				return err
			}
			cmd.Args[i] = arg[:j] + in.Name() + arg[j+3+len(ext):]
		} else if j := strings.Index(arg, "$out"); j != -1 {
			var err error
			ext := cmdArgExtension.FindString(arg[j+4:])
			if out, err = ioutil.TempFile("", "minify-out-*"+ext); err != nil {
				return err
			}
			cmd.Args[i] = arg[:j] + out.Name() + arg[j+4+len(ext):]
		}
	}

	if in == nil {
		cmd.Stdin = r
	} else if _, err := io.Copy(in, r); err != nil {
		return err
	}
	if out == nil {
		cmd.Stdout = w
	} else {
		defer io.Copy(w, out)
	}
	stderr := &bytes.Buffer{}
	cmd.Stderr = stderr

	err := cmd.Run()
	if _, ok := err.(*exec.ExitError); ok {
		if stderr.Len() != 0 {
			err = fmt.Errorf("%s", stderr.String())
		}
		err = fmt.Errorf("command %s failed: %w", cmd.Path, err)
	}
	return err
}

////////////////////////////////////////////////////////////////

// M holds a map of mimetype => function to allow recursive minifier calls of the minifier functions.
type M struct {
	mutex   sync.RWMutex
	literal map[string]Minifier
	pattern []patternMinifier

	URL *url.URL
}

// New returns a new M.
func New() *M {
	return &M{
		sync.RWMutex{},
		map[string]Minifier{},
		[]patternMinifier{},
		nil,
	}
}

// Add adds a minifier to the mimetype => function map (unsafe for concurrent use).
func (m *M) Add(mimetype string, minifier Minifier) {
	m.mutex.Lock()
	m.literal[mimetype] = minifier
	m.mutex.Unlock()
}

// AddFunc adds a minify function to the mimetype => function map (unsafe for concurrent use).
func (m *M) AddFunc(mimetype string, minifier MinifierFunc) {
	m.mutex.Lock()
	m.literal[mimetype] = minifier
	m.mutex.Unlock()
}

// AddRegexp adds a minifier to the mimetype => function map (unsafe for concurrent use).
func (m *M) AddRegexp(pattern *regexp.Regexp, minifier Minifier) {
	m.mutex.Lock()
	m.pattern = append(m.pattern, patternMinifier{pattern, minifier})
	m.mutex.Unlock()
}

// AddFuncRegexp adds a minify function to the mimetype => function map (unsafe for concurrent use).
func (m *M) AddFuncRegexp(pattern *regexp.Regexp, minifier MinifierFunc) {
	m.mutex.Lock()
	m.pattern = append(m.pattern, patternMinifier{pattern, minifier})
	m.mutex.Unlock()
}

// AddCmd adds a minify function to the mimetype => function map (unsafe for concurrent use) that executes a command to process the minification.
// It allows the use of external tools like ClosureCompiler, UglifyCSS, etc. for a specific mimetype.
func (m *M) AddCmd(mimetype string, cmd *exec.Cmd) {
	m.mutex.Lock()
	m.literal[mimetype] = &cmdMinifier{cmd}
	m.mutex.Unlock()
}

// AddCmdRegexp adds a minify function to the mimetype => function map (unsafe for concurrent use) that executes a command to process the minification.
// It allows the use of external tools like ClosureCompiler, UglifyCSS, etc. for a specific mimetype regular expression.
func (m *M) AddCmdRegexp(pattern *regexp.Regexp, cmd *exec.Cmd) {
	m.mutex.Lock()
	m.pattern = append(m.pattern, patternMinifier{pattern, &cmdMinifier{cmd}})
	m.mutex.Unlock()
}

// Match returns the pattern and minifier that gets matched with the mediatype.
// It returns nil when no matching minifier exists.
// It has the same matching algorithm as Minify.
func (m *M) Match(mediatype string) (string, map[string]string, MinifierFunc) {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	mimetype, params := parse.Mediatype([]byte(mediatype))
	if minifier, ok := m.literal[string(mimetype)]; ok { // string conversion is optimized away
		return string(mimetype), params, minifier.Minify
	}

	for _, minifier := range m.pattern {
		if minifier.pattern.Match(mimetype) {
			return minifier.pattern.String(), params, minifier.Minify
		}
	}
	return string(mimetype), params, nil
}

// Minify minifies the content of a Reader and writes it to a Writer (safe for concurrent use).
// An error is returned when no such mimetype exists (ErrNotExist) or when an error occurred in the minifier function.
// Mediatype may take the form of 'text/plain', 'text/*', '*/*' or 'text/plain; charset=UTF-8; version=2.0'.
func (m *M) Minify(mediatype string, w io.Writer, r io.Reader) error {
	mimetype, params := parse.Mediatype([]byte(mediatype))
	return m.MinifyMimetype(mimetype, w, r, params)
}

// MinifyMimetype minifies the content of a Reader and writes it to a Writer (safe for concurrent use).
// It is a lower level version of Minify and requires the mediatype to be split up into mimetype and parameters.
// It is mostly used internally by minifiers because it is faster (no need to convert a byte-slice to string and vice versa).
func (m *M) MinifyMimetype(mimetype []byte, w io.Writer, r io.Reader, params map[string]string) error {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	if minifier, ok := m.literal[string(mimetype)]; ok { // string conversion is optimized away
		return minifier.Minify(m, w, r, params)
	}
	for _, minifier := range m.pattern {
		if minifier.pattern.Match(mimetype) {
			return minifier.Minify(m, w, r, params)
		}
	}
	return ErrNotExist
}

// Bytes minifies an array of bytes (safe for concurrent use). When an error occurs it return the original array and the error.
// It returns an error when no such mimetype exists (ErrNotExist) or any error occurred in the minifier function.
func (m *M) Bytes(mediatype string, v []byte) ([]byte, error) {
	out := buffer.NewWriter(make([]byte, 0, len(v)))
	if err := m.Minify(mediatype, out, buffer.NewReader(v)); err != nil {
		return v, err
	}
	return out.Bytes(), nil
}

// String minifies a string (safe for concurrent use). When an error occurs it return the original string and the error.
// It returns an error when no such mimetype exists (ErrNotExist) or any error occurred in the minifier function.
func (m *M) String(mediatype string, v string) (string, error) {
	out := buffer.NewWriter(make([]byte, 0, len(v)))
	if err := m.Minify(mediatype, out, buffer.NewReader([]byte(v))); err != nil {
		return v, err
	}
	return string(out.Bytes()), nil
}

// Reader wraps a Reader interface and minifies the stream.
// Errors from the minifier are returned by the reader.
func (m *M) Reader(mediatype string, r io.Reader) io.Reader {
	pr, pw := io.Pipe()
	go func() {
		if err := m.Minify(mediatype, pw, r); err != nil {
			pw.CloseWithError(err)
		} else {
			pw.Close()
		}
	}()
	return pr
}

// writer makes sure that errors from the minifier are passed down through Close (can be blocking).
type writer struct {
	pw     *io.PipeWriter
	wg     sync.WaitGroup
	err    error
	closed bool
}

// Write intercepts any writes to the writer.
func (w *writer) Write(b []byte) (int, error) {
	if w.closed {
		return 0, ErrClosedWriter
	}
	n, err := w.pw.Write(b)
	if w.err != nil {
		err = w.err
	}
	return n, err
}

// Close must be called when writing has finished. It returns the error from the minifier.
func (w *writer) Close() error {
	if !w.closed {
		w.pw.Close()
		w.wg.Wait()
		w.closed = true
	}
	return w.err
}

// Writer wraps a Writer interface and minifies the stream.
// Errors from the minifier are returned by Close on the writer.
// The writer must be closed explicitly.
func (m *M) Writer(mediatype string, w io.Writer) *writer {
	pr, pw := io.Pipe()
	mw := &writer{pw, sync.WaitGroup{}, nil, false}
	mw.wg.Add(1)
	go func() {
		defer mw.wg.Done()

		if err := m.Minify(mediatype, w, pr); err != nil {
			mw.err = err
		}
		pr.Close()
	}()
	return mw
}

// responseWriter wraps an http.ResponseWriter and makes sure that errors from the minifier are passed down through Close (can be blocking).
// All writes to the response writer are intercepted and minified on the fly.
// http.ResponseWriter loses all functionality such as Pusher, Hijacker, Flusher, ...
type responseWriter struct {
	http.ResponseWriter

	writer    *writer
	m         *M
	mediatype string
}

// WriteHeader intercepts any header writes and removes the Content-Length header.
func (w *responseWriter) WriteHeader(status int) {
	w.ResponseWriter.Header().Del("Content-Length")
	w.ResponseWriter.WriteHeader(status)
}

// Write intercepts any writes to the response writer.
// The first write will extract the Content-Type as the mediatype. Otherwise it falls back to the RequestURI extension.
func (w *responseWriter) Write(b []byte) (int, error) {
	if w.writer == nil {
		// first write
		if mediatype := w.ResponseWriter.Header().Get("Content-Type"); mediatype != "" {
			w.mediatype = mediatype
		}
		w.writer = w.m.Writer(w.mediatype, w.ResponseWriter)
	}
	return w.writer.Write(b)
}

// Close must be called when writing has finished. It returns the error from the minifier.
func (w *responseWriter) Close() error {
	if w.writer != nil {
		return w.writer.Close()
	}
	return nil
}

// ResponseWriter minifies any writes to the http.ResponseWriter.
// http.ResponseWriter loses all functionality such as Pusher, Hijacker, Flusher, ...
// Minification might be slower than just sending the original file! Caching is advised.
func (m *M) ResponseWriter(w http.ResponseWriter, r *http.Request) *responseWriter {
	mediatype := mime.TypeByExtension(path.Ext(r.RequestURI))
	return &responseWriter{w, nil, m, mediatype}
}

// Middleware provides a middleware function that minifies content on the fly by intercepting writes to http.ResponseWriter.
// http.ResponseWriter loses all functionality such as Pusher, Hijacker, Flusher, ...
// Minification might be slower than just sending the original file! Caching is advised.
func (m *M) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mw := m.ResponseWriter(w, r)
		next.ServeHTTP(mw, r)
		mw.Close()
	})
}

// MiddlewareWithError provides a middleware function that minifies content on the fly by intercepting writes to http.ResponseWriter. The error function allows handling minification errors.
// http.ResponseWriter loses all functionality such as Pusher, Hijacker, Flusher, ...
// Minification might be slower than just sending the original file! Caching is advised.
func (m *M) MiddlewareWithError(next http.Handler, errorFunc func(w http.ResponseWriter, r *http.Request, err error)) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mw := m.ResponseWriter(w, r)
		next.ServeHTTP(mw, r)
		if err := mw.Close(); err != nil {
			errorFunc(w, r, err)
			return
		}
	})
}
<commit_msg>Change minify.writer.Writer return type to io.WriteCloser, fixes #524<commit_after>// Package minify relates MIME type to minifiers. Several minifiers are provided in the subpackages.
package minify

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path"
	"regexp"
	"strings"
	"sync"

	"github.com/tdewolff/parse/v2"
	"github.com/tdewolff/parse/v2/buffer"
)

// Warning is used to report usage warnings such as using a deprecated feature
var Warning = log.New(os.Stderr, "WARNING: ", 0)

// ErrNotExist is returned when no minifier exists for a given mimetype.
var ErrNotExist = errors.New("minifier does not exist for mimetype")

// ErrClosedWriter is returned when writing to a closed writer.
var ErrClosedWriter = errors.New("write on closed writer")

////////////////////////////////////////////////////////////////

// MinifierFunc is a function that implements Minifer.
type MinifierFunc func(*M, io.Writer, io.Reader, map[string]string) error

// Minify calls f(m, w, r, params)
func (f MinifierFunc) Minify(m *M, w io.Writer, r io.Reader, params map[string]string) error {
	return f(m, w, r, params)
}

// Minifier is the interface for minifiers.
// The *M parameter is used for minifying embedded resources, such as JS within HTML.
type Minifier interface {
	Minify(*M, io.Writer, io.Reader, map[string]string) error
}

////////////////////////////////////////////////////////////////

type patternMinifier struct {
	pattern *regexp.Regexp
	Minifier
}

type cmdMinifier struct {
	cmd *exec.Cmd
}

var cmdArgExtension = regexp.MustCompile(`^\.[0-9a-zA-Z]+`)

func (c *cmdMinifier) Minify(_ *M, w io.Writer, r io.Reader, _ map[string]string) error {
	cmd := &exec.Cmd{}
	*cmd = *c.cmd // concurrency safety

	var in, out *os.File
	for i, arg := range cmd.Args {
		if j := strings.Index(arg, "$in"); j != -1 {
			var err error
			ext := cmdArgExtension.FindString(arg[j+3:])
			if in, err = ioutil.TempFile("", "minify-in-*"+ext); err != nil {
				return err
			}
			cmd.Args[i] = arg[:j] + in.Name() + arg[j+3+len(ext):]
		} else if j := strings.Index(arg, "$out"); j != -1 {
			var err error
			ext := cmdArgExtension.FindString(arg[j+4:])
			if out, err = ioutil.TempFile("", "minify-out-*"+ext); err != nil {
				return err
			}
			cmd.Args[i] = arg[:j] + out.Name() + arg[j+4+len(ext):]
		}
	}

	if in == nil {
		cmd.Stdin = r
	} else if _, err := io.Copy(in, r); err != nil {
		return err
	}
	if out == nil {
		cmd.Stdout = w
	} else {
		defer io.Copy(w, out)
	}
	stderr := &bytes.Buffer{}
	cmd.Stderr = stderr

	err := cmd.Run()
	if _, ok := err.(*exec.ExitError); ok {
		if stderr.Len() != 0 {
			err = fmt.Errorf("%s", stderr.String())
		}
		err = fmt.Errorf("command %s failed: %w", cmd.Path, err)
	}
	return err
}

////////////////////////////////////////////////////////////////

// M holds a map of mimetype => function to allow recursive minifier calls of the minifier functions.
type M struct {
	mutex   sync.RWMutex
	literal map[string]Minifier
	pattern []patternMinifier

	URL *url.URL
}

// New returns a new M.
func New() *M {
	return &M{
		sync.RWMutex{},
		map[string]Minifier{},
		[]patternMinifier{},
		nil,
	}
}

// Add adds a minifier to the mimetype => function map (unsafe for concurrent use).
func (m *M) Add(mimetype string, minifier Minifier) {
	m.mutex.Lock()
	m.literal[mimetype] = minifier
	m.mutex.Unlock()
}

// AddFunc adds a minify function to the mimetype => function map (unsafe for concurrent use).
func (m *M) AddFunc(mimetype string, minifier MinifierFunc) {
	m.mutex.Lock()
	m.literal[mimetype] = minifier
	m.mutex.Unlock()
}

// AddRegexp adds a minifier to the mimetype => function map (unsafe for concurrent use).
func (m *M) AddRegexp(pattern *regexp.Regexp, minifier Minifier) {
	m.mutex.Lock()
	m.pattern = append(m.pattern, patternMinifier{pattern, minifier})
	m.mutex.Unlock()
}

// AddFuncRegexp adds a minify function to the mimetype => function map (unsafe for concurrent use).
func (m *M) AddFuncRegexp(pattern *regexp.Regexp, minifier MinifierFunc) {
	m.mutex.Lock()
	m.pattern = append(m.pattern, patternMinifier{pattern, minifier})
	m.mutex.Unlock()
}

// AddCmd adds a minify function to the mimetype => function map (unsafe for concurrent use) that executes a command to process the minification.
// It allows the use of external tools like ClosureCompiler, UglifyCSS, etc. for a specific mimetype.
func (m *M) AddCmd(mimetype string, cmd *exec.Cmd) {
	m.mutex.Lock()
	m.literal[mimetype] = &cmdMinifier{cmd}
	m.mutex.Unlock()
}

// AddCmdRegexp adds a minify function to the mimetype => function map (unsafe for concurrent use) that executes a command to process the minification.
// It allows the use of external tools like ClosureCompiler, UglifyCSS, etc. for a specific mimetype regular expression.
func (m *M) AddCmdRegexp(pattern *regexp.Regexp, cmd *exec.Cmd) {
	m.mutex.Lock()
	m.pattern = append(m.pattern, patternMinifier{pattern, &cmdMinifier{cmd}})
	m.mutex.Unlock()
}

// Match returns the pattern and minifier that gets matched with the mediatype.
// It returns nil when no matching minifier exists.
// It has the same matching algorithm as Minify.
func (m *M) Match(mediatype string) (string, map[string]string, MinifierFunc) {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	mimetype, params := parse.Mediatype([]byte(mediatype))
	if minifier, ok := m.literal[string(mimetype)]; ok { // string conversion is optimized away
		return string(mimetype), params, minifier.Minify
	}

	for _, minifier := range m.pattern {
		if minifier.pattern.Match(mimetype) {
			return minifier.pattern.String(), params, minifier.Minify
		}
	}
	return string(mimetype), params, nil
}

// Minify minifies the content of a Reader and writes it to a Writer (safe for concurrent use).
// An error is returned when no such mimetype exists (ErrNotExist) or when an error occurred in the minifier function.
// Mediatype may take the form of 'text/plain', 'text/*', '*/*' or 'text/plain; charset=UTF-8; version=2.0'.
func (m *M) Minify(mediatype string, w io.Writer, r io.Reader) error {
	mimetype, params := parse.Mediatype([]byte(mediatype))
	return m.MinifyMimetype(mimetype, w, r, params)
}

// MinifyMimetype minifies the content of a Reader and writes it to a Writer (safe for concurrent use).
// It is a lower level version of Minify and requires the mediatype to be split up into mimetype and parameters.
// It is mostly used internally by minifiers because it is faster (no need to convert a byte-slice to string and vice versa).
func (m *M) MinifyMimetype(mimetype []byte, w io.Writer, r io.Reader, params map[string]string) error {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	if minifier, ok := m.literal[string(mimetype)]; ok { // string conversion is optimized away
		return minifier.Minify(m, w, r, params)
	}
	for _, minifier := range m.pattern {
		if minifier.pattern.Match(mimetype) {
			return minifier.Minify(m, w, r, params)
		}
	}
	return ErrNotExist
}

// Bytes minifies an array of bytes (safe for concurrent use). When an error occurs it return the original array and the error.
// It returns an error when no such mimetype exists (ErrNotExist) or any error occurred in the minifier function.
func (m *M) Bytes(mediatype string, v []byte) ([]byte, error) {
	out := buffer.NewWriter(make([]byte, 0, len(v)))
	if err := m.Minify(mediatype, out, buffer.NewReader(v)); err != nil {
		return v, err
	}
	return out.Bytes(), nil
}

// String minifies a string (safe for concurrent use). When an error occurs it return the original string and the error.
// It returns an error when no such mimetype exists (ErrNotExist) or any error occurred in the minifier function.
func (m *M) String(mediatype string, v string) (string, error) {
	out := buffer.NewWriter(make([]byte, 0, len(v)))
	if err := m.Minify(mediatype, out, buffer.NewReader([]byte(v))); err != nil {
		return v, err
	}
	return string(out.Bytes()), nil
}

// Reader wraps a Reader interface and minifies the stream.
// Errors from the minifier are returned by the reader.
func (m *M) Reader(mediatype string, r io.Reader) io.Reader {
	pr, pw := io.Pipe()
	go func() {
		if err := m.Minify(mediatype, pw, r); err != nil {
			pw.CloseWithError(err)
		} else {
			pw.Close()
		}
	}()
	return pr
}

// writer makes sure that errors from the minifier are passed down through Close (can be blocking).
type writer struct {
	pw     *io.PipeWriter
	wg     sync.WaitGroup
	err    error
	closed bool
}

// Write intercepts any writes to the writer.
func (w *writer) Write(b []byte) (int, error) {
	if w.closed {
		return 0, ErrClosedWriter
	}
	n, err := w.pw.Write(b)
	if w.err != nil {
		err = w.err
	}
	return n, err
}

// Close must be called when writing has finished. It returns the error from the minifier.
func (w *writer) Close() error {
	if !w.closed {
		w.pw.Close()
		w.wg.Wait()
		w.closed = true
	}
	return w.err
}

// Writer wraps a Writer interface and minifies the stream.
// Errors from the minifier are returned by Close on the writer.
// The writer must be closed explicitly.
func (m *M) Writer(mediatype string, w io.Writer) io.WriteCloser {
	pr, pw := io.Pipe()
	mw := &writer{pw, sync.WaitGroup{}, nil, false}
	mw.wg.Add(1)
	go func() {
		defer mw.wg.Done()

		if err := m.Minify(mediatype, w, pr); err != nil {
			mw.err = err
		}
		pr.Close()
	}()
	return mw
}

// responseWriter wraps an http.ResponseWriter and makes sure that errors from the minifier are passed down through Close (can be blocking).
// All writes to the response writer are intercepted and minified on the fly.
// http.ResponseWriter loses all functionality such as Pusher, Hijacker, Flusher, ...
type responseWriter struct {
	http.ResponseWriter

	writer    *writer
	m         *M
	mediatype string
}

// WriteHeader intercepts any header writes and removes the Content-Length header.
func (w *responseWriter) WriteHeader(status int) {
	w.ResponseWriter.Header().Del("Content-Length")
	w.ResponseWriter.WriteHeader(status)
}

// Write intercepts any writes to the response writer.
// The first write will extract the Content-Type as the mediatype. Otherwise it falls back to the RequestURI extension.
func (w *responseWriter) Write(b []byte) (int, error) {
	if w.writer == nil {
		// first write
		if mediatype := w.ResponseWriter.Header().Get("Content-Type"); mediatype != "" {
			w.mediatype = mediatype
		}
		w.writer = w.m.Writer(w.mediatype, w.ResponseWriter).(*writer)
	}
	return w.writer.Write(b)
}

// Close must be called when writing has finished. It returns the error from the minifier.
func (w *responseWriter) Close() error {
	if w.writer != nil {
		return w.writer.Close()
	}
	return nil
}

// ResponseWriter minifies any writes to the http.ResponseWriter.
// http.ResponseWriter loses all functionality such as Pusher, Hijacker, Flusher, ...
// Minification might be slower than just sending the original file! Caching is advised.
func (m *M) ResponseWriter(w http.ResponseWriter, r *http.Request) *responseWriter {
	mediatype := mime.TypeByExtension(path.Ext(r.RequestURI))
	return &responseWriter{w, nil, m, mediatype}
}

// Middleware provides a middleware function that minifies content on the fly by intercepting writes to http.ResponseWriter.
// http.ResponseWriter loses all functionality such as Pusher, Hijacker, Flusher, ...
// Minification might be slower than just sending the original file! Caching is advised.
func (m *M) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mw := m.ResponseWriter(w, r)
		next.ServeHTTP(mw, r)
		mw.Close()
	})
}

// MiddlewareWithError provides a middleware function that minifies content on the fly by intercepting writes to http.ResponseWriter. The error function allows handling minification errors.
// http.ResponseWriter loses all functionality such as Pusher, Hijacker, Flusher, ...
// Minification might be slower than just sending the original file! Caching is advised.
func (m *M) MiddlewareWithError(next http.Handler, errorFunc func(w http.ResponseWriter, r *http.Request, err error)) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mw := m.ResponseWriter(w, r)
		next.ServeHTTP(mw, r)
		if err := mw.Close(); err != nil {
			errorFunc(w, r, err)
			return
		}
	})
}
<|endoftext|>


<commit_before>// Copyright 2018 Authors of Cilium
//
// 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 service

import (
	"encoding/json"
	"path"

	"github.com/cilium/cilium/pkg/kvstore"
	"github.com/cilium/cilium/pkg/loadbalancer"
)

var (
	// ServiceStorePrefix is the kvstore prefix of the shared store
	//
	// WARNING - STABLE API: Changing the structure or values of this will
	// break backwards compatibility
	ServiceStorePrefix = path.Join(kvstore.BaseKeyPrefix, "state", "services", "v1")
)

// PortConfiguration is the L4 port configuration of a frontend or backend. The
// map is indexed by the name of the port and the value constains the L4 port
// and protocol.
type PortConfiguration map[string]*loadbalancer.L4Addr

// DeepEquals returns true if both PortConfigurations are identical
func (p PortConfiguration) DeepEquals(o PortConfiguration) bool {
	if len(p) != len(o) {
		return false
	}

	for portName1, port1 := range p {
		port2, ok := o[portName1]

		if !ok || !port1.Equals(port2) {
			return false
		}
	}

	return true
}

// ClusterService is the definition of a service in a cluster
//
// WARNING - STABLE API: Any change to this structure must be done in a
// backwards compatible way.
type ClusterService struct {
	// Cluster is the cluster name the service is configured in
	Cluster string `json:"cluster"`

	// Namespace is the cluster namespace the service is configured in
	Namespace string `json:"namespace"`

	// Name is the name of the service. It must be unique within the
	// namespace of the cluster
	Name string `json:"name"`

	// Frontends is a map indexed by the frontend IP address
	Frontends map[string]PortConfiguration `json:"frontends"`

	// Backends is is map indexed by the backend IP address
	Backends map[string]PortConfiguration `json:"backends"`

	// Labels are the labels of the service
	Labels map[string]string `json:"labels"`

	// Selector is the label selector used to select backends
	Selector map[string]string `json:"selector"`
}

func (s *ClusterService) String() string {
	return s.Cluster + "/" + s.Namespace + "/" + s.Name
}

// NamespaceServiceName returns the namespace and service name
func (s *ClusterService) NamespaceServiceName() string {
	return s.Namespace + "/" + s.Name
}

// GetKeyName returns the kvstore key to be used for the global service
func (s *ClusterService) GetKeyName() string {
	// WARNING - STABLE API: Changing the structure of the key may break
	// backwards compatibility
	return path.Join(s.Cluster, s.Namespace, s.Name)
}

// Marshal returns the global service object as JSON byte slice
func (s *ClusterService) Marshal() ([]byte, error) {
	return json.Marshal(s)
}

// Unmarshal parses the JSON byte slice and updates the global service receiver
func (s *ClusterService) Unmarshal(data []byte) error {
	return json.Unmarshal(data, s)
}

// NewClusterService returns a new cluster service definition
func NewClusterService(name, namespace string) ClusterService {
	return ClusterService{
		Name:      name,
		Namespace: namespace,
		Frontends: map[string]PortConfiguration{},
		Backends:  map[string]PortConfiguration{},
		Labels:    map[string]string{},
		Selector:  map[string]string{},
	}
}
<commit_msg>service: Fix Unmarshal() of GlobalService<commit_after>// Copyright 2018 Authors of Cilium
//
// 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 service

import (
	"encoding/json"
	"path"

	"github.com/cilium/cilium/pkg/kvstore"
	"github.com/cilium/cilium/pkg/loadbalancer"
)

var (
	// ServiceStorePrefix is the kvstore prefix of the shared store
	//
	// WARNING - STABLE API: Changing the structure or values of this will
	// break backwards compatibility
	ServiceStorePrefix = path.Join(kvstore.BaseKeyPrefix, "state", "services", "v1")
)

// PortConfiguration is the L4 port configuration of a frontend or backend. The
// map is indexed by the name of the port and the value constains the L4 port
// and protocol.
type PortConfiguration map[string]*loadbalancer.L4Addr

// DeepEquals returns true if both PortConfigurations are identical
func (p PortConfiguration) DeepEquals(o PortConfiguration) bool {
	if len(p) != len(o) {
		return false
	}

	for portName1, port1 := range p {
		port2, ok := o[portName1]

		if !ok || !port1.Equals(port2) {
			return false
		}
	}

	return true
}

// ClusterService is the definition of a service in a cluster
//
// WARNING - STABLE API: Any change to this structure must be done in a
// backwards compatible way.
type ClusterService struct {
	// Cluster is the cluster name the service is configured in
	Cluster string `json:"cluster"`

	// Namespace is the cluster namespace the service is configured in
	Namespace string `json:"namespace"`

	// Name is the name of the service. It must be unique within the
	// namespace of the cluster
	Name string `json:"name"`

	// Frontends is a map indexed by the frontend IP address
	Frontends map[string]PortConfiguration `json:"frontends"`

	// Backends is is map indexed by the backend IP address
	Backends map[string]PortConfiguration `json:"backends"`

	// Labels are the labels of the service
	Labels map[string]string `json:"labels"`

	// Selector is the label selector used to select backends
	Selector map[string]string `json:"selector"`
}

func (s *ClusterService) String() string {
	return s.Cluster + "/" + s.Namespace + "/" + s.Name
}

// NamespaceServiceName returns the namespace and service name
func (s *ClusterService) NamespaceServiceName() string {
	return s.Namespace + "/" + s.Name
}

// GetKeyName returns the kvstore key to be used for the global service
func (s *ClusterService) GetKeyName() string {
	// WARNING - STABLE API: Changing the structure of the key may break
	// backwards compatibility
	return path.Join(s.Cluster, s.Namespace, s.Name)
}

// Marshal returns the global service object as JSON byte slice
func (s *ClusterService) Marshal() ([]byte, error) {
	return json.Marshal(s)
}

// Unmarshal parses the JSON byte slice and updates the global service receiver
func (s *ClusterService) Unmarshal(data []byte) error {
	newService := NewClusterService("", "")

	if err := json.Unmarshal(data, &newService); err != nil {
		return err
	}

	*s = newService

	return nil
}

// NewClusterService returns a new cluster service definition
func NewClusterService(name, namespace string) ClusterService {
	return ClusterService{
		Name:      name,
		Namespace: namespace,
		Frontends: map[string]PortConfiguration{},
		Backends:  map[string]PortConfiguration{},
		Labels:    map[string]string{},
		Selector:  map[string]string{},
	}
}
<|endoftext|>


<commit_before>package store

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	level "github.com/go-kit/kit/log/experimental_level"
	"github.com/prometheus/client_golang/prometheus"

	"github.com/oklog/prototype/pkg/cluster"
	"github.com/oklog/prototype/pkg/ingest"
)

// Consumer reads segments from the ingesters, and replicates merged segments to
// the rest of the cluster. It's implemented as a state machine: gather
// segments, replicate, commit, and repeat. All failures invalidate the entire
// batch.
type Consumer struct {
	peer               *cluster.Peer
	client             *http.Client
	segmentTargetSize  int64
	replicationFactor  int
	gatherErrors       int                 // heuristic to move out of gather state
	pending            map[string][]string // ingester: segment IDs
	active             *bytes.Buffer       // merged pending segments
	activeSince        time.Time           // active segment has been "open" since this time
	stop               chan chan struct{}
	consumedSegments   prometheus.Counter
	consumedBytes      prometheus.Counter
	replicatedSegments prometheus.Counter
	replicatedBytes    prometheus.Counter
	logger             log.Logger
}

// NewConsumer creates a consumer.
// Don't forget to Run it.
func NewConsumer(
	peer *cluster.Peer,
	client *http.Client,
	segmentTargetSize int64,
	replicationFactor int,
	consumedSegments, consumedBytes prometheus.Counter,
	replicatedSegments, replicatedBytes prometheus.Counter,
	logger log.Logger,
) *Consumer {
	return &Consumer{
		peer:               peer,
		client:             client,
		segmentTargetSize:  segmentTargetSize,
		replicationFactor:  replicationFactor,
		gatherErrors:       0,
		pending:            map[string][]string{},
		active:             &bytes.Buffer{},
		activeSince:        time.Time{},
		stop:               make(chan chan struct{}),
		consumedSegments:   consumedSegments,
		consumedBytes:      consumedBytes,
		replicatedSegments: replicatedSegments,
		replicatedBytes:    replicatedBytes,
		logger:             logger,
	}
}

// Run consumes segments from ingest nodes, and replicates them to the cluster.
// Run returns when Stop is invoked.
func (c *Consumer) Run() {
	step := time.Tick(100 * time.Millisecond)
	state := c.gather
	for {
		select {
		case <-step:
			state = state()

		case q := <-c.stop:
			c.fail() // any outstanding segments
			close(q)
			return
		}
	}
}

// Stop the consumer from consuming.
func (c *Consumer) Stop() {
	q := make(chan struct{})
	c.stop <- q
	<-q
}

type stateFn func() stateFn

func (c *Consumer) gather() stateFn {
	var (
		base = log.NewContext(c.logger).With("state", "gather")
		warn = level.Warn(base)
	)

	// A naïve way to break out of the gather loop in atypical conditions.
	// TODO(pb): this obviously needs more thought and consideration
	instances := c.peer.Current(cluster.PeerTypeIngest)
	if c.gatherErrors > 0 && c.gatherErrors > 2*len(instances) {
		if c.active.Len() <= 0 {
			// We didn't successfully consume any segments.
			// Nothing to do but reset and try again.
			c.gatherErrors = 0
			return c.gather
		}
		// We consumed some segment, at least.
		// Press forward to persistence.
		return c.replicate
	}
	if len(instances) == 0 {
		return c.gather // maybe some will come back later
	}
	if want, have := c.replicationFactor, len(c.peer.Current(cluster.PeerTypeStore)); have < want {
		// Don't gather if we can't replicate.
		// Better to queue up on the ingesters.
		warn.Log("replication_factor", want, "available_peers", have, "err", "replication currently impossible")
		time.Sleep(time.Second)
		c.gatherErrors++
		return c.gather
	}

	// More typical exit clauses.
	const maxAge = time.Second // TODO(pb): parameterize?
	var (
		tooBig = int64(c.active.Len()) > c.segmentTargetSize
		tooOld = !c.activeSince.IsZero() && time.Now().Sub(c.activeSince) > maxAge
	)
	if tooBig || tooOld {
		return c.replicate
	}

	// Get the oldest segment ID from a random ingester.
	instance := instances[rand.Intn(len(instances))]
	nextResp, err := c.client.Get(fmt.Sprintf("http://%s/ingest%s", instance, ingest.APIPathNext))
	if err != nil {
		warn.Log("ingester", instance, "during", ingest.APIPathNext, "err", err)
		c.gatherErrors++
		return c.gather
	}
	defer nextResp.Body.Close()
	nextRespBody, err := ioutil.ReadAll(nextResp.Body)
	if err != nil {
		warn.Log("ingester", instance, "during", ingest.APIPathNext, "err", err)
		c.gatherErrors++
		return c.gather
	}
	nextID := strings.TrimSpace(string(nextRespBody))
	if nextResp.StatusCode == http.StatusNotFound {
		// Normal, when the ingester has no more segments to give right now.
		c.gatherErrors++ // after enough of these errors, we should replicate
		return c.gather
	}
	if nextResp.StatusCode != http.StatusOK {
		warn.Log("ingester", instance, "during", ingest.APIPathNext, "returned", nextResp.Status)
		c.gatherErrors++
		return c.gather
	}

	// Mark the segment ID as pending.
	// From this point forward, we must either commit or fail the segment.
	// If we do neither, it will eventually time out, but we should be nice.
	c.pending[instance] = append(c.pending[instance], nextID)

	// Read the segment.
	readResp, err := c.client.Get(fmt.Sprintf("http://%s/ingest%s?id=%s", instance, ingest.APIPathRead, nextID))
	if err != nil {
		// Reading failed, so we can't possibly commit the segment.
		// The simplest thing to do now is to fail everything.
		// TODO(pb): this could be improved i.e. made more granular
		warn.Log("ingester", instance, "during", ingest.APIPathRead, "err", err)
		c.gatherErrors++
		return c.fail // fail everything
	}
	defer readResp.Body.Close()
	if readResp.StatusCode != http.StatusOK {
		warn.Log("ingester", instance, "during", ingest.APIPathRead, "returned", readResp.Status)
		c.gatherErrors++
		return c.fail // fail everything, same as above
	}

	// Merge the segment into our active segment.
	var cw countingWriter
	if _, _, _, err := mergeRecords(c.active, io.TeeReader(readResp.Body, &cw)); err != nil {
		warn.Log("ingester", instance, "during", "mergeRecords", "err", err)
		c.gatherErrors++
		return c.fail // fail everything, same as above
	}
	if c.activeSince.IsZero() {
		c.activeSince = time.Now()
	}

	// Repeat!
	c.consumedSegments.Inc()
	c.consumedBytes.Add(float64(cw.n))
	return c.gather
}

func (c *Consumer) replicate() stateFn {
	var (
		base = log.NewContext(c.logger).With("state", "replicate")
		warn = level.Warn(base)
	)

	// Replicate the segment to the cluster.
	var (
		peers      = c.peer.Current(cluster.PeerTypeStore)
		indices    = rand.Perm(len(peers))
		replicated = 0
	)
	if want, have := c.replicationFactor, len(peers); have < want {
		warn.Log("replication_factor", want, "available_peers", have, "err", "replication currently impossible")
		return c.fail // can't do anything here
	}
	for i := 0; i < len(indices) && replicated < c.replicationFactor; i++ {
		var (
			index    = indices[i]
			target   = peers[index]
			uri      = fmt.Sprintf("http://%s/store%s", target, APIPathReplicate)
			bodyType = "application/binary"
			body     = bytes.NewReader(c.active.Bytes())
		)
		resp, err := c.client.Post(uri, bodyType, body)
		if err != nil {
			warn.Log("target", target, "during", APIPathReplicate, "err", err)
			continue // we'll try another one
		}
		resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			warn.Log("target", target, "during", APIPathReplicate, "got", resp.Status)
			continue // we'll try another one
		}
		replicated++
	}
	if replicated < c.replicationFactor {
		warn.Log("err", "failed to fully replicate", "want", c.replicationFactor, "have", replicated)
		return c.fail // harsh, but OK
	}

	// All good!
	c.replicatedSegments.Inc()
	c.replicatedBytes.Add(float64(c.active.Len()))
	return c.commit
}

func (c *Consumer) commit() stateFn {
	return c.resetVia("commit")
}

func (c *Consumer) fail() stateFn {
	return c.resetVia("failed")
}

func (c *Consumer) resetVia(commitOrFailed string) stateFn {
	var (
		base = log.NewContext(c.logger).With("state", commitOrFailed)
		warn = level.Warn(base)
	)

	// If commits fail, the segment may be re-replicated; that's OK.
	// If fails fail, the segment will eventually time-out; that's also OK.
	// So we have best-effort semantics, just log the error and move on.
	var wg sync.WaitGroup
	for instance, ids := range c.pending {
		wg.Add(len(ids))
		for _, id := range ids {
			go func(id string) {
				defer wg.Done()
				uri := fmt.Sprintf("http://%s/ingest/%s?id=%s", instance, commitOrFailed, id)
				resp, err := c.client.Post(uri, "text/plain", nil)
				if err != nil {
					warn.Log("instance", instance, "during", "POST", "uri", uri, "err", err)
					return
				}
				resp.Body.Close()
				if resp.StatusCode != http.StatusOK {
					warn.Log("instance", instance, "during", "POST", "status", resp.Status)
					return
				}
			}(id)
		}
	}
	wg.Wait()

	// Reset various pending things.
	c.gatherErrors = 0
	c.pending = map[string][]string{}
	c.active.Reset()
	c.activeSince = time.Time{}

	// Back to the beginning.
	return c.gather
}

type countingWriter struct{ n int64 }

func (cw *countingWriter) Write(p []byte) (int, error) {
	cw.n += int64(len(p))
	return len(p), nil
}
<commit_msg>Something weird<commit_after>package store

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	level "github.com/go-kit/kit/log/experimental_level"
	"github.com/prometheus/client_golang/prometheus"

	"github.com/oklog/prototype/pkg/cluster"
	"github.com/oklog/prototype/pkg/ingest"
)

// Consumer reads segments from the ingesters, and replicates merged segments to
// the rest of the cluster. It's implemented as a state machine: gather
// segments, replicate, commit, and repeat. All failures invalidate the entire
// batch.
type Consumer struct {
	peer               *cluster.Peer
	client             *http.Client
	segmentTargetSize  int64
	replicationFactor  int
	gatherErrors       int                 // heuristic to move out of gather state
	pending            map[string][]string // ingester: segment IDs
	active             *bytes.Buffer       // merged pending segments
	activeSince        time.Time           // active segment has been "open" since this time
	stop               chan chan struct{}
	consumedSegments   prometheus.Counter
	consumedBytes      prometheus.Counter
	replicatedSegments prometheus.Counter
	replicatedBytes    prometheus.Counter
	logger             log.Logger
}

// NewConsumer creates a consumer.
// Don't forget to Run it.
func NewConsumer(
	peer *cluster.Peer,
	client *http.Client,
	segmentTargetSize int64,
	replicationFactor int,
	consumedSegments, consumedBytes prometheus.Counter,
	replicatedSegments, replicatedBytes prometheus.Counter,
	logger log.Logger,
) *Consumer {
	return &Consumer{
		peer:               peer,
		client:             client,
		segmentTargetSize:  segmentTargetSize,
		replicationFactor:  replicationFactor,
		gatherErrors:       0,
		pending:            map[string][]string{},
		active:             &bytes.Buffer{},
		activeSince:        time.Time{},
		stop:               make(chan chan struct{}),
		consumedSegments:   consumedSegments,
		consumedBytes:      consumedBytes,
		replicatedSegments: replicatedSegments,
		replicatedBytes:    replicatedBytes,
		logger:             logger,
	}
}

// Run consumes segments from ingest nodes, and replicates them to the cluster.
// Run returns when Stop is invoked.
func (c *Consumer) Run() {
	step := time.Tick(100 * time.Millisecond)
	state := c.gather
	for {
		select {
		case <-step:
			state = state()

		case q := <-c.stop:
			c.fail() // any outstanding segments
			close(q)
			return
		}
	}
}

// Stop the consumer from consuming.
func (c *Consumer) Stop() {
	q := make(chan struct{})
	c.stop <- q
	<-q
}

type stateFn func() stateFn

func (c *Consumer) gather() stateFn {
	var (
		base = log.NewContext(c.logger).With("state", "gather")
		warn = level.Warn(base)
	)

	// A naïve way to break out of the gather loop in atypical conditions.
	// TODO(pb): this obviously needs more thought and consideration
	instances := c.peer.Current(cluster.PeerTypeIngest)
	if c.gatherErrors > 0 && c.gatherErrors > 2*len(instances) {
		if c.active.Len() <= 0 {
			// We didn't successfully consume any segments.
			// Nothing to do but reset and try again.
			c.gatherErrors = 0
			return c.gather
		}
		// We consumed some segment, at least.
		// Press forward to persistence.
		return c.replicate
	}
	if len(instances) == 0 {
		return c.gather // maybe some will come back later
	}
	if want, have := c.replicationFactor, len(c.peer.Current(cluster.PeerTypeStore)); have < want {
		// Don't gather if we can't replicate.
		// Better to queue up on the ingesters.
		warn.Log("replication_factor", want, "available_peers", have, "err", "replication currently impossible")
		time.Sleep(time.Second)
		c.gatherErrors++
		return c.gather
	}

	// More typical exit clauses.
	const maxAge = time.Second // TODO(pb): parameterize?
	var (
		tooBig = int64(c.active.Len()) > c.segmentTargetSize
		tooOld = !c.activeSince.IsZero() && time.Now().Sub(c.activeSince) > maxAge
	)
	if tooBig || tooOld {
		return c.replicate
	}

	// Get the oldest segment ID from a random ingester.
	instance := instances[rand.Intn(len(instances))]
	nextResp, err := c.client.Get(fmt.Sprintf("http://%s/ingest%s", instance, ingest.APIPathNext))
	if err != nil {
		warn.Log("ingester", instance, "during", ingest.APIPathNext, "err", err)
		c.gatherErrors++
		return c.gather
	}
	defer nextResp.Body.Close()
	nextRespBody, err := ioutil.ReadAll(nextResp.Body)
	if err != nil {
		warn.Log("ingester", instance, "during", ingest.APIPathNext, "err", err)
		c.gatherErrors++
		return c.gather
	}
	nextID := strings.TrimSpace(string(nextRespBody))
	if nextResp.StatusCode == http.StatusNotFound {
		// Normal, when the ingester has no more segments to give right now.
		c.gatherErrors++ // after enough of these errors, we should replicate
		return c.gather
	}
	if nextResp.StatusCode != http.StatusOK {
		warn.Log("ingester", instance, "during", ingest.APIPathNext, "returned", nextResp.Status)
		c.gatherErrors++
		return c.gather
	}

	// Mark the segment ID as pending.
	// From this point forward, we must either commit or fail the segment.
	// If we do neither, it will eventually time out, but we should be nice.
	c.pending[instance] = append(c.pending[instance], nextID)

	// Read the segment.
	readResp, err := c.client.Get(fmt.Sprintf("http://%s/ingest%s?id=%s", instance, ingest.APIPathRead, nextID))
	if err != nil {
		// Reading failed, so we can't possibly commit the segment.
		// The simplest thing to do now is to fail everything.
		// TODO(pb): this could be improved i.e. made more granular
		warn.Log("ingester", instance, "during", ingest.APIPathRead, "err", err)
		c.gatherErrors++
		return c.fail // fail everything
	}
	defer readResp.Body.Close()
	if readResp.StatusCode != http.StatusOK {
		warn.Log("ingester", instance, "during", ingest.APIPathRead, "returned", readResp.Status)
		c.gatherErrors++
		return c.fail // fail everything, same as above
	}

	// Merge the segment into our active segment.
	var cw countingWriter
	if _, _, _, err := mergeRecords(c.active, io.TeeReader(readResp.Body, &cw)); err != nil {
		warn.Log("ingester", instance, "during", "mergeRecords", "err", err)
		c.gatherErrors++
		return c.fail // fail everything, same as above
	}
	if c.activeSince.IsZero() {
		c.activeSince = time.Now()
	}

	// Repeat!
	c.consumedSegments.Inc()
	c.consumedBytes.Add(float64(cw.n))
	return c.gather
}

func (c *Consumer) replicate() stateFn {
	var (
		base = log.NewContext(c.logger).With("state", "replicate")
		warn = level.Warn(base)
	)

	// Replicate the segment to the cluster.
	var (
		peers      = c.peer.Current(cluster.PeerTypeStore)
		indices    = rand.Perm(len(peers))
		replicated = 0
	)
	if want, have := c.replicationFactor, len(peers); have < want {
		warn.Log("replication_factor", want, "available_peers", have, "err", "replication currently impossible")
		return c.fail // can't do anything here
	}
	for i := 0; i < len(indices) && replicated < c.replicationFactor; i++ {
		var (
			index    = indices[i]
			target   = peers[index]
			uri      = fmt.Sprintf("http://%s/store%s", target, APIPathReplicate)
			bodyType = "application/binary"
			body     = bytes.NewReader(c.active.Bytes())
		)
		resp, err := c.client.Post(uri, bodyType, body)
		if err != nil {
			warn.Log("target", target, "during", APIPathReplicate, "err", err)
			continue // we'll try another one
		}
		resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			warn.Log("target", target, "during", APIPathReplicate, "got", resp.Status)
			continue // we'll try another one
		}
		replicated++
	}
	if replicated < c.replicationFactor {
		warn.Log("err", "failed to fully replicate", "want", c.replicationFactor, "have", replicated)
		return c.fail // harsh, but OK
	}

	// All good!
	c.replicatedSegments.Inc()
	c.replicatedBytes.Add(float64(c.active.Len()))
	return c.commit
}

func (c *Consumer) commit() stateFn {
	return c.resetVia("commit")
}

func (c *Consumer) fail() stateFn {
	return c.resetVia("failed")
}

func (c *Consumer) resetVia(commitOrFailed string) stateFn {
	var (
		base = log.NewContext(c.logger).With("state", commitOrFailed)
		warn = level.Warn(base)
	)

	// If commits fail, the segment may be re-replicated; that's OK.
	// If fails fail, the segment will eventually time-out; that's also OK.
	// So we have best-effort semantics, just log the error and move on.
	var wg sync.WaitGroup
	for instance, ids := range c.pending {
		wg.Add(len(ids))
		for _, id := range ids {
			go func(id string) {
				defer wg.Done()
				uri := fmt.Sprintf("http://%s/ingest/%s?id=%s", instance, commitOrFailed, id)
				resp, err := c.client.Post(uri, "text/plain", nil)
				if err != nil {
					warn.Log("instance", instance, "during", "POST", "uri", uri, "err", err)
					return
				}
				resp.Body.Close()
				if resp.StatusCode != http.StatusOK {
					warn.Log("instance", instance, "during", "POST", "uri", uri, "status", resp.Status)
					return
				}
			}(id)
		}
	}
	wg.Wait()

	// Reset various pending things.
	c.gatherErrors = 0
	c.pending = map[string][]string{}
	c.active.Reset()
	c.activeSince = time.Time{}

	// Back to the beginning.
	return c.gather
}

type countingWriter struct{ n int64 }

func (cw *countingWriter) Write(p []byte) (int, error) {
	cw.n += int64(len(p))
	return len(p), nil
}
<|endoftext|>


<commit_before>/*
Copyright 2016 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 cert

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	cryptorand "crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"errors"
	"fmt"
	"net"

	"k8s.io/kubernetes/pkg/apis/certificates"
)

// ParseCSR extracts the CSR from the API object and decodes it.
func ParseCSR(obj *certificates.CertificateSigningRequest) (*x509.CertificateRequest, error) {
	// extract PEM from request object
	pemBytes := obj.Spec.Request
	block, _ := pem.Decode(pemBytes)
	if block == nil || block.Type != "CERTIFICATE REQUEST" {
		return nil, errors.New("PEM block type must be CERTIFICATE REQUEST")
	}
	csr, err := x509.ParseCertificateRequest(block.Bytes)
	if err != nil {
		return nil, err
	}
	return csr, nil
}

// MakeCSR generates a PEM-encoded CSR using the supplied private key, subject, and SANs.
// privateKey must be a *ecdsa.PrivateKey or *rsa.PrivateKey.
func MakeCSR(privateKey interface{}, subject *pkix.Name, dnsSANs []string, ipSANs []net.IP) (csr []byte, err error) {
	var sigType x509.SignatureAlgorithm

	switch privateKey := privateKey.(type) {
	case *ecdsa.PrivateKey:
		switch privateKey.Curve {
		case elliptic.P224(), elliptic.P256():
			sigType = x509.ECDSAWithSHA256
		case elliptic.P384():
			sigType = x509.ECDSAWithSHA384
		case elliptic.P521():
			sigType = x509.ECDSAWithSHA512
		default:
			return nil, fmt.Errorf("unknown elliptic curve: %v", privateKey.Curve)
		}
	case *rsa.PrivateKey:
		keySize := privateKey.N.BitLen()
		switch {
		case keySize >= 4096:
			sigType = x509.SHA512WithRSA
		case keySize >= 3072:
			sigType = x509.SHA384WithRSA
		default:
			sigType = x509.SHA256WithRSA
		}

	default:
		return nil, fmt.Errorf("unsupported key type: %T", privateKey)
	}

	template := &x509.CertificateRequest{
		Subject:            *subject,
		SignatureAlgorithm: sigType,
		DNSNames:           dnsSANs,
		IPAddresses:        ipSANs,
	}

	csr, err = x509.CreateCertificateRequest(cryptorand.Reader, template, privateKey)
	if err != nil {
		return nil, err
	}

	csrPemBlock := &pem.Block{
		Type:  "CERTIFICATE REQUEST",
		Bytes: csr,
	}

	return pem.EncodeToMemory(csrPemBlock), nil
}
<commit_msg>Remove duplicated ECDHE key handling<commit_after>/*
Copyright 2016 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 cert

import (
	cryptorand "crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"errors"
	"net"

	"k8s.io/kubernetes/pkg/apis/certificates"
)

// ParseCSR extracts the CSR from the API object and decodes it.
func ParseCSR(obj *certificates.CertificateSigningRequest) (*x509.CertificateRequest, error) {
	// extract PEM from request object
	pemBytes := obj.Spec.Request
	block, _ := pem.Decode(pemBytes)
	if block == nil || block.Type != "CERTIFICATE REQUEST" {
		return nil, errors.New("PEM block type must be CERTIFICATE REQUEST")
	}
	csr, err := x509.ParseCertificateRequest(block.Bytes)
	if err != nil {
		return nil, err
	}
	return csr, nil
}

// MakeCSR generates a PEM-encoded CSR using the supplied private key, subject, and SANs.
// All key types that are implemented via crypto.Signer are supported (This includes *rsa.PrivateKey and *ecdsa.PrivateKey.)
func MakeCSR(privateKey interface{}, subject *pkix.Name, dnsSANs []string, ipSANs []net.IP) (csr []byte, err error) {
	// Customize the signature for RSA keys, depending on the key size
	var sigType x509.SignatureAlgorithm
	if privateKey, ok := privateKey.(*rsa.PrivateKey); ok {
		keySize := privateKey.N.BitLen()
		switch {
		case keySize >= 4096:
			sigType = x509.SHA512WithRSA
		case keySize >= 3072:
			sigType = x509.SHA384WithRSA
		default:
			sigType = x509.SHA256WithRSA
		}
	}

	template := &x509.CertificateRequest{
		Subject:            *subject,
		SignatureAlgorithm: sigType,
		DNSNames:           dnsSANs,
		IPAddresses:        ipSANs,
	}

	csr, err = x509.CreateCertificateRequest(cryptorand.Reader, template, privateKey)
	if err != nil {
		return nil, err
	}

	csrPemBlock := &pem.Block{
		Type:  "CERTIFICATE REQUEST",
		Bytes: csr,
	}

	return pem.EncodeToMemory(csrPemBlock), nil
}
<|endoftext|>


<commit_before>package the_platinum_searcher

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"

	"github.com/monochromegane/conflag"
	"github.com/monochromegane/go-home"
	"github.com/monochromegane/terminal"
)

const version = "2.0.3"

const (
	ExitCodeOK = iota
	ExitCodeError
)

var opts Option

type PlatinumSearcher struct {
	Out, Err io.Writer
}

func (p PlatinumSearcher) Run(args []string) int {

	parser := newOptionParser(&opts)

	conflag.LongHyphen = true
	conflag.BoolValue = false
	for _, c := range []string{filepath.Join(home.Dir(), ".ptconfig.toml"), ".ptconfig.toml"} {
		if args, err := conflag.ArgsFrom(c); err == nil {
			parser.ParseArgs(args)
		}
	}

	args, err := parser.ParseArgs(args)
	if err != nil {
		return ExitCodeError
	}

	if opts.Version {
		fmt.Printf("pt version %s\n", version)
		return ExitCodeOK
	}

	if len(args) == 0 && !opts.SearchOption.EnableFilesWithRegexp {
		parser.WriteHelp(p.Err)
		return ExitCodeError
	}

	if !terminal.IsTerminal(os.Stdout) {
		opts.OutputOption.EnableColor = false
		opts.OutputOption.EnableGroup = false
	}

	if p.givenStdin() {
		opts.SearchOption.SearchStream = true
	}

	if opts.SearchOption.EnableFilesWithRegexp {
		args = append([]string{""}, args...)
	}

	if opts.OutputOption.Count {
		opts.OutputOption.Before = 0
		opts.OutputOption.After = 0
		opts.OutputOption.Context = 0
	}

	search := search{
		roots: p.rootsFrom(args),
		out:   p.Out,
	}
	if err = search.start(p.patternFrom(args)); err != nil {
		fmt.Fprintf(p.Err, "%s\n", err)
		return ExitCodeError
	}
	return ExitCodeOK
}

func (p PlatinumSearcher) patternFrom(args []string) string {
	return args[0]
}

func (p PlatinumSearcher) rootsFrom(args []string) []string {
	if len(args) > 1 {
		return args[1:]
	} else {
		return []string{"."}
	}
}

func (p PlatinumSearcher) givenStdin() bool {
	fi, err := os.Stdin.Stat()
	if runtime.GOOS == "windows" {
		if err == nil {
			return true
		}
	} else {
		if err != nil {
			return false
		}

		mode := fi.Mode()
		if (mode&os.ModeNamedPipe != 0) || mode.IsRegular() {
			return true
		}
	}
	return false
}
<commit_msg>Bumped version to 2.0.4.<commit_after>package the_platinum_searcher

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"

	"github.com/monochromegane/conflag"
	"github.com/monochromegane/go-home"
	"github.com/monochromegane/terminal"
)

const version = "2.0.4"

const (
	ExitCodeOK = iota
	ExitCodeError
)

var opts Option

type PlatinumSearcher struct {
	Out, Err io.Writer
}

func (p PlatinumSearcher) Run(args []string) int {

	parser := newOptionParser(&opts)

	conflag.LongHyphen = true
	conflag.BoolValue = false
	for _, c := range []string{filepath.Join(home.Dir(), ".ptconfig.toml"), ".ptconfig.toml"} {
		if args, err := conflag.ArgsFrom(c); err == nil {
			parser.ParseArgs(args)
		}
	}

	args, err := parser.ParseArgs(args)
	if err != nil {
		return ExitCodeError
	}

	if opts.Version {
		fmt.Printf("pt version %s\n", version)
		return ExitCodeOK
	}

	if len(args) == 0 && !opts.SearchOption.EnableFilesWithRegexp {
		parser.WriteHelp(p.Err)
		return ExitCodeError
	}

	if !terminal.IsTerminal(os.Stdout) {
		opts.OutputOption.EnableColor = false
		opts.OutputOption.EnableGroup = false
	}

	if p.givenStdin() {
		opts.SearchOption.SearchStream = true
	}

	if opts.SearchOption.EnableFilesWithRegexp {
		args = append([]string{""}, args...)
	}

	if opts.OutputOption.Count {
		opts.OutputOption.Before = 0
		opts.OutputOption.After = 0
		opts.OutputOption.Context = 0
	}

	search := search{
		roots: p.rootsFrom(args),
		out:   p.Out,
	}
	if err = search.start(p.patternFrom(args)); err != nil {
		fmt.Fprintf(p.Err, "%s\n", err)
		return ExitCodeError
	}
	return ExitCodeOK
}

func (p PlatinumSearcher) patternFrom(args []string) string {
	return args[0]
}

func (p PlatinumSearcher) rootsFrom(args []string) []string {
	if len(args) > 1 {
		return args[1:]
	} else {
		return []string{"."}
	}
}

func (p PlatinumSearcher) givenStdin() bool {
	fi, err := os.Stdin.Stat()
	if runtime.GOOS == "windows" {
		if err == nil {
			return true
		}
	} else {
		if err != nil {
			return false
		}

		mode := fi.Mode()
		if (mode&os.ModeNamedPipe != 0) || mode.IsRegular() {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package cast

func Uint32(v interface{}) (uint32, error) {

	switch value := v.(type) {
	case uint8:
		return uint32(value), nil
	case uint16:
		return uint32(value), nil
	case uint32:
		return uint32(value), nil
	case uint32er:
		return value.Uint32()
	default:
		return 0, internalCannotCastComplainer{expectedType:"uint32", actualType:typeof(value)}
	}
}

// MustUint32 is like Uint32, expect panic()s on an error.
func MustUint32(v interface{}) uint32 {

	x, err := Uint32(v)
	if nil != err {
		panic(err)
	}

	return x
}

type uint32er interface {
	Uint32() (uint32, error)
}
<commit_msg>added more docs. and improved cast.Uint32()<commit_after>package cast

// Uint32 will return an int32 when `v` is of type uint32, uint16, uint8 or has a method:
//
//	type interface {
//		Uint32() (uint32, error)
//	}
//
// ... that returns successfully, or has a method:
//
//	type interface {
//		Uint16() (uint16, error)
//	}
//
// ... that returns successfully, or has a method:
//
//	type interface {
//		Uint8() (uint8, error)
//	}
//
// ... that returns successfully.
//
// Else it will return an error.
func Uint32(v interface{}) (uint32, error) {

	switch value := v.(type) {
	case uint32:
		return uint32(value), nil
	case uint32er:
		return value.Uint32()
	case uint16:
		return uint32(value), nil
	case uint16er:
		return func()(uint32, error){
			casted, err := value.Uint16()
			if nil != err {
				return 0, err
			}
			return uint32(casted), nil
		}()
	case uint8:
		return uint32(value), nil
	case uint8er:
		return func()(uint32, error){
			casted, err := value.Uint8()
			if nil != err {
				return 0, err
			}
			return uint32(casted), nil
		}()
	default:
		return 0, internalCannotCastComplainer{expectedType:"uint32", actualType:typeof(value)}
	}
}

// MustUint32 is like Uint32, expect panic()s on an error.
func MustUint32(v interface{}) uint32 {

	x, err := Uint32(v)
	if nil != err {
		panic(err)
	}

	return x
}

type uint32er interface {
	Uint32() (uint32, error)

}
<|endoftext|>


<commit_before>// Copyright 2015 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package upload

import (
	"errors"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"time"
)

// Upload用于处理文件上传
type Upload struct {
	dir     string   // 上传文件保存的路径根目录
	maxSize int64    // 允许的最大文件大小，以byte为单位
	role    string   // 文件命名方式
	exts    []string // 允许的扩展名
}

// 声明一个Upload对象。
// dir 上传文件的保存目录;
// maxSize 允许上传文件的最大尺寸，单位为byte；
// role 文件命名规则，格式可参考time.Format()参数；
// exts 允许的扩展名，若为空，将不允许任何文件上传。
func New(dir string, maxSize int64, role string, exts ...string) *Upload {
	es := make([]string, 0, len(exts))
	for _, ext := range exts {
		if ext[0] != '.' {
			es = append(es, "."+ext)
			continue
		}
		es = append(es, ext)
	}

	return &Upload{
		dir:     dir,
		maxSize: maxSize,
		role:    role,
		exts:    es,
	}
}

// 判断扩展名是否符合要求。
func (u *Upload) checkExt(ext string) bool {
	if len(ext) == 0 { // 没有扩展名，一律过滤
		return false
	}

	// 是否为允许的扩展名
	for _, e := range u.exts {
		if e == ext {
			return true
		}
	}
	return false
}

// 检测文件大小是否符合要求。
func (u *Upload) checkSize(file multipart.File) (bool, error) {
	var size int64

	switch f := file.(type) {
	case stater:
		stat, err := f.Stat()
		if err != nil {
			return false, err
		}
		size = stat.Size()
	case sizer:
		size = f.Size()
	default:
		return false, errors.New("上传文件时发生未知的错误")
	}

	return size <= u.maxSize, nil
}

// 设置水印，file为水印文件的路径，或是在isText为true时，file为水印的文字。
func (u *Upload) SetWaterMark(file string, isText bool) {
	// TODO
}

// 招行上传的操作。会检测上传文件是否符合要求，只要有一个文件不符合，就会中断上传。
// 返回的是相对于u.dir目录的文件名列表。
func (u *Upload) Do(field string, w *http.ResponseWriter, r *http.Request) ([]string, error) {
	r.ParseMultipartForm(32 << 20)
	heads := r.MultipartForm.File[field]
	ret := make([]string, len(heads))

	for _, head := range heads {
		file, err := head.Open()
		if err != nil {
			return nil, err
		}

		ext := filepath.Ext(head.Filename)
		if !u.checkExt(ext) {
			return nil, errors.New("包含无效的文件类型")
		}

		ok, err := u.checkSize(file)
		if err != nil {
			return nil, err
		}
		if !ok {
			return nil, errors.New("超过最大的文件大小")
		}

		path := time.Now().Format(u.role) + ext
		ret = append(ret, path)
		f, err := os.Create(u.dir + path)
		if err != nil {
			return nil, err
		}

		io.Copy(f, file)

		f.Close()
		file.Close() // for的最后关闭file
	}

	return ret, nil
}
<commit_msg>为New的dir参数添加检测代码 1.确保dir参数的最后一个字符为/符号； 2.检测dir目录是否真实存在，若不存在，会尝试创建该目录；<commit_after>// Copyright 2015 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package upload

import (
	"errors"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"time"
)

// 创建文件的默认权限，比如Upload.dir若不存在，会使用此权限创建目录。
const defaultMode os.FileMode = 0660

// Upload用于处理文件上传
type Upload struct {
	dir     string   // 上传文件保存的路径根目录
	maxSize int64    // 允许的最大文件大小，以byte为单位
	role    string   // 文件命名方式
	exts    []string // 允许的扩展名
}

// 声明一个Upload对象。
// dir 上传文件的保存目录，若目录不存在，则会尝试创建;
// maxSize 允许上传文件的最大尺寸，单位为byte；
// role 文件命名规则，格式可参考time.Format()参数；
// exts 允许的扩展名，若为空，将不允许任何文件上传。
func New(dir string, maxSize int64, role string, exts ...string) (*Upload, error) {
	// 确保所有的后缀名都是以.作为开始符号的。
	es := make([]string, 0, len(exts))
	for _, ext := range exts {
		if ext[0] != '.' {
			es = append(es, "."+ext)
			continue
		}
		es = append(es, ext)
	}

	// 确保dir最后一个字符为目录分隔符。
	last := dir[len(dir)-1]
	if last != '/' && last != filepath.Separator {
		dir = dir + string(filepath.Separator)
	}

	// 确保dir目录存在，若不存在则会尝试创建。
	stat, err := os.Stat(dir)
	if err != nil && !os.IsExist(err) {
		if !os.IsNotExist(err) {
			return nil, err
		}

		// 尝试创建目录
		if err = os.MkdirAll(dir, 0660); err != nil {
			return nil, err
		}

		// 创建目录成功，重新获取状态
		if stat, err = os.Stat(dir); err != nil {
			return nil, err
		}
	}
	if !stat.IsDir() {
		return nil, errors.New("dir不是一个目录")
	}

	return &Upload{
		dir:     dir,
		maxSize: maxSize,
		role:    role,
		exts:    es,
	}, nil
}

// 判断扩展名是否符合要求。
func (u *Upload) checkExt(ext string) bool {
	if len(ext) == 0 { // 没有扩展名，一律过滤
		return false
	}

	// 是否为允许的扩展名
	for _, e := range u.exts {
		if e == ext {
			return true
		}
	}
	return false
}

// 检测文件大小是否符合要求。
func (u *Upload) checkSize(file multipart.File) (bool, error) {
	var size int64

	switch f := file.(type) {
	case stater:
		stat, err := f.Stat()
		if err != nil {
			return false, err
		}
		size = stat.Size()
	case sizer:
		size = f.Size()
	default:
		return false, errors.New("上传文件时发生未知的错误")
	}

	return size <= u.maxSize, nil
}

// 设置水印，file为水印文件的路径，或是在isText为true时，file为水印的文字。
func (u *Upload) SetWaterMark(file string, isText bool) {
	// TODO
}

// 招行上传的操作。会检测上传文件是否符合要求，只要有一个文件不符合，就会中断上传。
// 返回的是相对于u.dir目录的文件名列表。
func (u *Upload) Do(field string, w *http.ResponseWriter, r *http.Request) ([]string, error) {
	r.ParseMultipartForm(32 << 20)
	heads := r.MultipartForm.File[field]
	ret := make([]string, len(heads))

	for _, head := range heads {
		file, err := head.Open()
		if err != nil {
			return nil, err
		}

		ext := filepath.Ext(head.Filename)
		if !u.checkExt(ext) {
			return nil, errors.New("包含无效的文件类型")
		}

		ok, err := u.checkSize(file)
		if err != nil {
			return nil, err
		}
		if !ok {
			return nil, errors.New("超过最大的文件大小")
		}

		path := time.Now().Format(u.role) + ext
		ret = append(ret, path)
		f, err := os.Create(u.dir + path)
		if err != nil {
			return nil, err
		}

		io.Copy(f, file)

		f.Close()
		file.Close() // for的最后关闭file
	}

	return ret, nil
}
<|endoftext|>


<commit_before>package main

import (

	//	"encoding/json"
	"crypto/sha1"
	"encoding/base64"
	"fmt"
	"io"
	"path/filepath"

	"github.com/dchest/uniuri"
	//	"mime/multipart"
	"net/http"
	"os"
)

const (
	LENGTH    = 6
	PORT      = ":8080"
	DIRECTORY = "/tmp/"
	UPADDRESS = "http://localhost/"
)

type Result struct {
	URL  string `json:"url"`
	Name string `json:"name"`
	Hash string `json:"hash"`
	Size int64  `jason:"size"`
}

type Response struct {
	SUCCESS     bool     `json:"success"`
	ERRORCODE   int      `json:"errorcode,omitempty"`
	DESCRIPTION string   `json:"description,omitempty"`
	FILES       []Result `json:"files,omitempty"`
}

func generateName() string {
	name := uniuri.NewLen(LENGTH)
	return name
}
func check(err error) {
	if err != nil {
		fmt.Println(err)
		return
	}
}

func uploadHandler(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		io.WriteString(w, "Error")
	case "POST":

		reader, err := r.MultipartReader()

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		for {
			part, err := reader.NextPart()
			if err == io.EOF {
				break
			}

			if part.FileName() == "" {
				continue
			}
			s := generateName()
			extName := filepath.Ext(part.FileName())
			filename := s + extName
			dst, err := os.Create(DIRECTORY + filename)
			defer dst.Close()

			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}

			h := sha1.New()
			t := io.TeeReader(part, h)
			_, err = io.Copy(dst, t)

			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			hash := h.Sum(nil)
			sha1 := base64.URLEncoding.EncodeToString(hash)
			size := dst.Stat()
			res := Result{
				URL:  UPADDRESS + "/" + s,
				Name: part.FileName(),
				Hash: sha1,
				Size: size.Size(),
			}

			io.WriteString(w, filename+"\n")
		}
	}

}

func main() {
	http.HandleFunc("/upload.php", uploadHandler)
	err := http.ListenAndServe(PORT, nil)
	if err != nil {
		panic(err)
	}

}
<commit_msg>Fix struct response<commit_after>package main

import (

	//	"encoding/json"
	"crypto/sha1"
	"encoding/base64"
	"fmt"
	"io"
	"path/filepath"

	"github.com/dchest/uniuri"
	//	"mime/multipart"
	"net/http"
	"os"
)

const (
	LENGTH    = 6
	PORT      = ":8080"
	DIRECTORY = "/tmp/"
	UPADDRESS = "http://localhost"
)

type Result struct {
	URL  string `json:"url"`
	Name string `json:"name"`
	Hash string `json:"hash"`
	Size int64  `jason:"size"`
}

type Response struct {
	SUCCESS     bool     `json:"success"`
	ERRORCODE   int      `json:"errorcode,omitempty"`
	DESCRIPTION string   `json:"description,omitempty"`
	FILES       []Result `json:"files,omitempty"`
}

func generateName() string {
	name := uniuri.NewLen(LENGTH)
	return name
}
func check(err error) {
	if err != nil {
		fmt.Println(err)
		return
	}
}

func uploadHandler(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		io.WriteString(w, "Error")
	case "POST":

		reader, err := r.MultipartReader()

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		for {
			part, err := reader.NextPart()
			if err == io.EOF {
				break
			}

			if part.FileName() == "" {
				continue
			}
			s := generateName()
			extName := filepath.Ext(part.FileName())
			filename := s + extName
			dst, err := os.Create(DIRECTORY + filename)
			defer dst.Close()

			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}

			h := sha1.New()
			t := io.TeeReader(part, h)
			_, err = io.Copy(dst, t)

			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			hash := h.Sum(nil)
			sha1 := base64.URLEncoding.EncodeToString(hash)
			size, _ := dst.Stat()
			res := Result{
				URL:  UPADDRESS + "/" + s + extName,
				Name: part.FileName(),
				Hash: sha1,
				Size: size.Size(),
			}

			fmt.Println(res)

		}
	}

}

func main() {
	http.HandleFunc("/upload.php", uploadHandler)
	err := http.ListenAndServe(PORT, nil)
	if err != nil {
		panic(err)
	}

}
<|endoftext|>


<commit_before>// Copyright 2013, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package zktopo

import (
	"encoding/json"
	"fmt"
	"path"
	"sort"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/jscfg"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/zk"
	"golang.org/x/net/context"
	"launchpad.net/gozk/zookeeper"
)

// WatchSleepDuration is how many seconds interval to poll for in case
// the directory that contains a file to watch doesn't exist, or a watch
// is broken. It is exported so individual test and main programs
// can change it.
var WatchSleepDuration = 30 * time.Second

/*
This file contains the serving graph management code of zktopo.Server
*/
func zkPathForCell(cell string) string {
	return fmt.Sprintf("/zk/%v/vt/ns", cell)
}

func zkPathForVtKeyspace(cell, keyspace string) string {
	return path.Join(zkPathForCell(cell), keyspace)
}

func zkPathForVtShard(cell, keyspace, shard string) string {
	return path.Join(zkPathForVtKeyspace(cell, keyspace), shard)
}

func zkPathForVtName(cell, keyspace, shard string, tabletType topo.TabletType) string {
	return path.Join(zkPathForVtShard(cell, keyspace, shard), string(tabletType))
}

// GetSrvTabletTypesPerShard is part of the topo.Server interface
func (zkts *Server) GetSrvTabletTypesPerShard(ctx context.Context, cell, keyspace, shard string) ([]topo.TabletType, error) {
	zkSgShardPath := zkPathForVtShard(cell, keyspace, shard)
	children, _, err := zkts.zconn.Children(zkSgShardPath)
	if err != nil {
		if zookeeper.IsError(err, zookeeper.ZNONODE) {
			err = topo.ErrNoNode
		}
		return nil, err
	}
	result := make([]topo.TabletType, 0, len(children))
	for _, tt := range children {
		// these two are used for locking
		if tt == "action" || tt == "actionlog" {
			continue
		}
		result = append(result, topo.TabletType(tt))
	}
	return result, nil
}

// CreateEndPoints is part of the topo.Server interface
func (zkts *Server) CreateEndPoints(ctx context.Context, cell, keyspace, shard string, tabletType topo.TabletType, addrs *topo.EndPoints) error {
	path := zkPathForVtName(cell, keyspace, shard, tabletType)
	data := jscfg.ToJSON(addrs)

	// Create only if it doesn't exist.
	_, err := zk.CreateRecursive(zkts.zconn, path, data, 0, zookeeper.WorldACL(zookeeper.PERM_ALL))
	if zookeeper.IsError(err, zookeeper.ZNODEEXISTS) {
		err = topo.ErrNodeExists
	}
	return err
}

// UpdateEndPoints is part of the topo.Server interface
func (zkts *Server) UpdateEndPoints(ctx context.Context, cell, keyspace, shard string, tabletType topo.TabletType, addrs *topo.EndPoints, existingVersion int64) error {
	path := zkPathForVtName(cell, keyspace, shard, tabletType)
	data := jscfg.ToJSON(addrs)

	if existingVersion == -1 {
		// Update or create unconditionally.
		_, err := zk.CreateRecursive(zkts.zconn, path, data, 0, zookeeper.WorldACL(zookeeper.PERM_ALL))
		if err != nil {
			if zookeeper.IsError(err, zookeeper.ZNODEEXISTS) {
				// Node already exists - just stomp away. Multiple writers shouldn't be here.
				// We use RetryChange here because it won't update the node unnecessarily.
				f := func(oldValue string, oldStat zk.Stat) (string, error) {
					return data, nil
				}
				err = zkts.zconn.RetryChange(path, 0, zookeeper.WorldACL(zookeeper.PERM_ALL), f)
			}
		}
		return err
	}

	// Compare And Set
	_, err := zkts.zconn.Set(path, data, int(existingVersion))
	if err != nil {
		if zookeeper.IsError(err, zookeeper.ZBADVERSION) {
			err = topo.ErrBadVersion
		} else if zookeeper.IsError(err, zookeeper.ZNONODE) {
			err = topo.ErrNoNode
		}
	}
	return err
}

// GetEndPoints is part of the topo.Server interface
func (zkts *Server) GetEndPoints(ctx context.Context, cell, keyspace, shard string, tabletType topo.TabletType) (*topo.EndPoints, int64, error) {
	path := zkPathForVtName(cell, keyspace, shard, tabletType)
	data, stat, err := zkts.zconn.Get(path)
	if err != nil {
		if zookeeper.IsError(err, zookeeper.ZNONODE) {
			err = topo.ErrNoNode
		}
		return nil, 0, err
	}
	result := &topo.EndPoints{}
	if len(data) > 0 {
		if err := json.Unmarshal([]byte(data), result); err != nil {
			return nil, 0, fmt.Errorf("EndPoints unmarshal failed: %v %v", data, err)
		}
	}
	return result, int64(stat.Version()), nil
}

// DeleteEndPoints is part of the topo.Server interface
func (zkts *Server) DeleteEndPoints(ctx context.Context, cell, keyspace, shard string, tabletType topo.TabletType, existingVersion int64) error {
	path := zkPathForVtName(cell, keyspace, shard, tabletType)
	if err := zkts.zconn.Delete(path, int(existingVersion)); err != nil {
		switch {
		case zookeeper.IsError(err, zookeeper.ZNONODE):
			err = topo.ErrNoNode
		case zookeeper.IsError(err, zookeeper.ZBADVERSION):
			err = topo.ErrBadVersion
		}
		return err
	}
	return nil
}

// UpdateSrvShard is part of the topo.Server interface
func (zkts *Server) UpdateSrvShard(ctx context.Context, cell, keyspace, shard string, srvShard *topo.SrvShard) error {
	path := zkPathForVtShard(cell, keyspace, shard)
	data := jscfg.ToJSON(srvShard)
	_, err := zkts.zconn.Set(path, data, -1)
	return err
}

// GetSrvShard is part of the topo.Server interface
func (zkts *Server) GetSrvShard(ctx context.Context, cell, keyspace, shard string) (*topo.SrvShard, error) {
	path := zkPathForVtShard(cell, keyspace, shard)
	data, stat, err := zkts.zconn.Get(path)
	if err != nil {
		if zookeeper.IsError(err, zookeeper.ZNONODE) {
			err = topo.ErrNoNode
		}
		return nil, err
	}
	srvShard := topo.NewSrvShard(int64(stat.Version()))
	if len(data) > 0 {
		if err := json.Unmarshal([]byte(data), srvShard); err != nil {
			return nil, fmt.Errorf("SrvShard unmarshal failed: %v %v", data, err)
		}
	}
	return srvShard, nil
}

// DeleteSrvShard is part of the topo.Server interface
func (zkts *Server) DeleteSrvShard(ctx context.Context, cell, keyspace, shard string) error {
	path := zkPathForVtShard(cell, keyspace, shard)
	err := zkts.zconn.Delete(path, -1)
	if err != nil {
		if zookeeper.IsError(err, zookeeper.ZNONODE) {
			err = topo.ErrNoNode
		}
		return err
	}
	return nil
}

// UpdateSrvKeyspace is part of the topo.Server interface
func (zkts *Server) UpdateSrvKeyspace(ctx context.Context, cell, keyspace string, srvKeyspace *topo.SrvKeyspace) error {
	path := zkPathForVtKeyspace(cell, keyspace)
	data := jscfg.ToJSON(srvKeyspace)
	_, err := zkts.zconn.Set(path, data, -1)
	if zookeeper.IsError(err, zookeeper.ZNONODE) {
		_, err = zk.CreateRecursive(zkts.zconn, path, data, 0, zookeeper.WorldACL(zookeeper.PERM_ALL))
	}
	return err
}

// GetSrvKeyspace is part of the topo.Server interface
func (zkts *Server) GetSrvKeyspace(ctx context.Context, cell, keyspace string) (*topo.SrvKeyspace, error) {
	path := zkPathForVtKeyspace(cell, keyspace)
	data, stat, err := zkts.zconn.Get(path)
	if err != nil {
		if zookeeper.IsError(err, zookeeper.ZNONODE) {
			err = topo.ErrNoNode
		}
		return nil, err
	}
	srvKeyspace := topo.NewSrvKeyspace(int64(stat.Version()))
	if len(data) > 0 {
		if err := json.Unmarshal([]byte(data), srvKeyspace); err != nil {
			return nil, fmt.Errorf("SrvKeyspace unmarshal failed: %v %v", data, err)
		}
	}
	return srvKeyspace, nil
}

// GetSrvKeyspaceNames is part of the topo.Server interface
func (zkts *Server) GetSrvKeyspaceNames(ctx context.Context, cell string) ([]string, error) {
	children, _, err := zkts.zconn.Children(zkPathForCell(cell))
	if err != nil {
		if zookeeper.IsError(err, zookeeper.ZNONODE) {
			return nil, nil
		}
		return nil, err
	}

	sort.Strings(children)
	return children, nil
}

var errSkipUpdate = fmt.Errorf("skip update")

func (zkts *Server) updateTabletEndpoint(oldValue string, oldStat zk.Stat, addr *topo.EndPoint) (newValue string, err error) {
	if oldStat == nil {
		// The incoming object doesn't exist - we haven't been placed in the serving
		// graph yet, so don't update. Assume the next process that rebuilds the graph
		// will get the updated tablet location.
		return "", errSkipUpdate
	}

	var addrs *topo.EndPoints
	if oldValue != "" {
		addrs = &topo.EndPoints{}
		if len(oldValue) > 0 {
			if err := json.Unmarshal([]byte(oldValue), addrs); err != nil {
				return "", fmt.Errorf("EndPoints unmarshal failed: %v %v", oldValue, err)
			}
		}

		foundTablet := false
		for i, entry := range addrs.Entries {
			if entry.Uid == addr.Uid {
				foundTablet = true
				if !topo.EndPointEquality(&entry, addr) {
					addrs.Entries[i] = *addr
				}
				break
			}
		}

		if !foundTablet {
			addrs.Entries = append(addrs.Entries, *addr)
		}
	} else {
		addrs = topo.NewEndPoints()
		addrs.Entries = append(addrs.Entries, *addr)
	}
	return jscfg.ToJSON(addrs), nil
}

// WatchEndPoints is part of the topo.Server interface
func (zkts *Server) WatchEndPoints(ctx context.Context, cell, keyspace, shard string, tabletType topo.TabletType) (<-chan *topo.EndPoints, chan<- struct{}, error) {
	filePath := zkPathForVtName(cell, keyspace, shard, tabletType)

	notifications := make(chan *topo.EndPoints, 10)
	stopWatching := make(chan struct{})

	// waitOrInterrupted will return true if stopWatching is triggered
	waitOrInterrupted := func() bool {
		timer := time.After(WatchSleepDuration)
		select {
		case <-stopWatching:
			close(notifications)
			return true
		case <-timer:
		}
		return false
	}

	go func() {
		for {
			// set the watch
			data, _, watch, err := zkts.zconn.GetW(filePath)
			if err != nil {
				if zookeeper.IsError(err, zookeeper.ZNONODE) {
					// the parent directory doesn't exist
					notifications <- nil
				}

				log.Errorf("Cannot set watch on %v, waiting for %v to retry: %v", filePath, WatchSleepDuration, err)
				if waitOrInterrupted() {
					return
				}
				continue
			}

			// get the initial value, send it, or send nil if no
			// data
			var ep *topo.EndPoints
			sendIt := true
			if len(data) > 0 {
				ep = &topo.EndPoints{}
				if err := json.Unmarshal([]byte(data), ep); err != nil {
					log.Errorf("EndPoints unmarshal failed: %v %v", data, err)
					sendIt = false
				}
			}
			if sendIt {
				notifications <- ep
			}

			// now act on the watch
			select {
			case event, ok := <-watch:
				if !ok {
					log.Warningf("watch on %v was closed, waiting for %v to retry", filePath, WatchSleepDuration)
					if waitOrInterrupted() {
						return
					}
					continue
				}

				if !event.Ok() {
					log.Warningf("received a non-OK event for %v, waiting for %v to retry", filePath, WatchSleepDuration)
					if waitOrInterrupted() {
						return
					}
				}
			case <-stopWatching:
				// user is not interested any more
				close(notifications)
				return
			}
		}
	}()

	return notifications, stopWatching, nil

}
<commit_msg>zktopo UpdateSrvShard should create parent if necessary.<commit_after>// Copyright 2013, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package zktopo

import (
	"encoding/json"
	"fmt"
	"path"
	"sort"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/jscfg"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/zk"
	"golang.org/x/net/context"
	"launchpad.net/gozk/zookeeper"
)

// WatchSleepDuration is how many seconds interval to poll for in case
// the directory that contains a file to watch doesn't exist, or a watch
// is broken. It is exported so individual test and main programs
// can change it.
var WatchSleepDuration = 30 * time.Second

/*
This file contains the serving graph management code of zktopo.Server
*/
func zkPathForCell(cell string) string {
	return fmt.Sprintf("/zk/%v/vt/ns", cell)
}

func zkPathForVtKeyspace(cell, keyspace string) string {
	return path.Join(zkPathForCell(cell), keyspace)
}

func zkPathForVtShard(cell, keyspace, shard string) string {
	return path.Join(zkPathForVtKeyspace(cell, keyspace), shard)
}

func zkPathForVtName(cell, keyspace, shard string, tabletType topo.TabletType) string {
	return path.Join(zkPathForVtShard(cell, keyspace, shard), string(tabletType))
}

// GetSrvTabletTypesPerShard is part of the topo.Server interface
func (zkts *Server) GetSrvTabletTypesPerShard(ctx context.Context, cell, keyspace, shard string) ([]topo.TabletType, error) {
	zkSgShardPath := zkPathForVtShard(cell, keyspace, shard)
	children, _, err := zkts.zconn.Children(zkSgShardPath)
	if err != nil {
		if zookeeper.IsError(err, zookeeper.ZNONODE) {
			err = topo.ErrNoNode
		}
		return nil, err
	}
	result := make([]topo.TabletType, 0, len(children))
	for _, tt := range children {
		// these two are used for locking
		if tt == "action" || tt == "actionlog" {
			continue
		}
		result = append(result, topo.TabletType(tt))
	}
	return result, nil
}

// CreateEndPoints is part of the topo.Server interface
func (zkts *Server) CreateEndPoints(ctx context.Context, cell, keyspace, shard string, tabletType topo.TabletType, addrs *topo.EndPoints) error {
	path := zkPathForVtName(cell, keyspace, shard, tabletType)
	data := jscfg.ToJSON(addrs)

	// Create only if it doesn't exist.
	_, err := zk.CreateRecursive(zkts.zconn, path, data, 0, zookeeper.WorldACL(zookeeper.PERM_ALL))
	if zookeeper.IsError(err, zookeeper.ZNODEEXISTS) {
		err = topo.ErrNodeExists
	}
	return err
}

// UpdateEndPoints is part of the topo.Server interface
func (zkts *Server) UpdateEndPoints(ctx context.Context, cell, keyspace, shard string, tabletType topo.TabletType, addrs *topo.EndPoints, existingVersion int64) error {
	path := zkPathForVtName(cell, keyspace, shard, tabletType)
	data := jscfg.ToJSON(addrs)

	if existingVersion == -1 {
		// Update or create unconditionally.
		_, err := zk.CreateRecursive(zkts.zconn, path, data, 0, zookeeper.WorldACL(zookeeper.PERM_ALL))
		if err != nil {
			if zookeeper.IsError(err, zookeeper.ZNODEEXISTS) {
				// Node already exists - just stomp away. Multiple writers shouldn't be here.
				// We use RetryChange here because it won't update the node unnecessarily.
				f := func(oldValue string, oldStat zk.Stat) (string, error) {
					return data, nil
				}
				err = zkts.zconn.RetryChange(path, 0, zookeeper.WorldACL(zookeeper.PERM_ALL), f)
			}
		}
		return err
	}

	// Compare And Set
	_, err := zkts.zconn.Set(path, data, int(existingVersion))
	if err != nil {
		if zookeeper.IsError(err, zookeeper.ZBADVERSION) {
			err = topo.ErrBadVersion
		} else if zookeeper.IsError(err, zookeeper.ZNONODE) {
			err = topo.ErrNoNode
		}
	}
	return err
}

// GetEndPoints is part of the topo.Server interface
func (zkts *Server) GetEndPoints(ctx context.Context, cell, keyspace, shard string, tabletType topo.TabletType) (*topo.EndPoints, int64, error) {
	path := zkPathForVtName(cell, keyspace, shard, tabletType)
	data, stat, err := zkts.zconn.Get(path)
	if err != nil {
		if zookeeper.IsError(err, zookeeper.ZNONODE) {
			err = topo.ErrNoNode
		}
		return nil, 0, err
	}
	result := &topo.EndPoints{}
	if len(data) > 0 {
		if err := json.Unmarshal([]byte(data), result); err != nil {
			return nil, 0, fmt.Errorf("EndPoints unmarshal failed: %v %v", data, err)
		}
	}
	return result, int64(stat.Version()), nil
}

// DeleteEndPoints is part of the topo.Server interface
func (zkts *Server) DeleteEndPoints(ctx context.Context, cell, keyspace, shard string, tabletType topo.TabletType, existingVersion int64) error {
	path := zkPathForVtName(cell, keyspace, shard, tabletType)
	if err := zkts.zconn.Delete(path, int(existingVersion)); err != nil {
		switch {
		case zookeeper.IsError(err, zookeeper.ZNONODE):
			err = topo.ErrNoNode
		case zookeeper.IsError(err, zookeeper.ZBADVERSION):
			err = topo.ErrBadVersion
		}
		return err
	}
	return nil
}

// UpdateSrvShard is part of the topo.Server interface
func (zkts *Server) UpdateSrvShard(ctx context.Context, cell, keyspace, shard string, srvShard *topo.SrvShard) error {
	path := zkPathForVtShard(cell, keyspace, shard)
	data := jscfg.ToJSON(srvShard)

	// Update or create unconditionally.
	_, err := zk.CreateRecursive(zkts.zconn, path, data, 0, zookeeper.WorldACL(zookeeper.PERM_ALL))
	if err != nil {
		if zookeeper.IsError(err, zookeeper.ZNODEEXISTS) {
			// Node already exists - just stomp away. Multiple writers shouldn't be here.
			// We use RetryChange here because it won't update the node unnecessarily.
			f := func(oldValue string, oldStat zk.Stat) (string, error) {
				return data, nil
			}
			err = zkts.zconn.RetryChange(path, 0, zookeeper.WorldACL(zookeeper.PERM_ALL), f)
		}
	}
	return err
}

// GetSrvShard is part of the topo.Server interface
func (zkts *Server) GetSrvShard(ctx context.Context, cell, keyspace, shard string) (*topo.SrvShard, error) {
	path := zkPathForVtShard(cell, keyspace, shard)
	data, stat, err := zkts.zconn.Get(path)
	if err != nil {
		if zookeeper.IsError(err, zookeeper.ZNONODE) {
			err = topo.ErrNoNode
		}
		return nil, err
	}
	srvShard := topo.NewSrvShard(int64(stat.Version()))
	if len(data) > 0 {
		if err := json.Unmarshal([]byte(data), srvShard); err != nil {
			return nil, fmt.Errorf("SrvShard unmarshal failed: %v %v", data, err)
		}
	}
	return srvShard, nil
}

// DeleteSrvShard is part of the topo.Server interface
func (zkts *Server) DeleteSrvShard(ctx context.Context, cell, keyspace, shard string) error {
	path := zkPathForVtShard(cell, keyspace, shard)
	err := zkts.zconn.Delete(path, -1)
	if err != nil {
		if zookeeper.IsError(err, zookeeper.ZNONODE) {
			err = topo.ErrNoNode
		}
		return err
	}
	return nil
}

// UpdateSrvKeyspace is part of the topo.Server interface
func (zkts *Server) UpdateSrvKeyspace(ctx context.Context, cell, keyspace string, srvKeyspace *topo.SrvKeyspace) error {
	path := zkPathForVtKeyspace(cell, keyspace)
	data := jscfg.ToJSON(srvKeyspace)
	_, err := zkts.zconn.Set(path, data, -1)
	if zookeeper.IsError(err, zookeeper.ZNONODE) {
		_, err = zk.CreateRecursive(zkts.zconn, path, data, 0, zookeeper.WorldACL(zookeeper.PERM_ALL))
	}
	return err
}

// GetSrvKeyspace is part of the topo.Server interface
func (zkts *Server) GetSrvKeyspace(ctx context.Context, cell, keyspace string) (*topo.SrvKeyspace, error) {
	path := zkPathForVtKeyspace(cell, keyspace)
	data, stat, err := zkts.zconn.Get(path)
	if err != nil {
		if zookeeper.IsError(err, zookeeper.ZNONODE) {
			err = topo.ErrNoNode
		}
		return nil, err
	}
	srvKeyspace := topo.NewSrvKeyspace(int64(stat.Version()))
	if len(data) > 0 {
		if err := json.Unmarshal([]byte(data), srvKeyspace); err != nil {
			return nil, fmt.Errorf("SrvKeyspace unmarshal failed: %v %v", data, err)
		}
	}
	return srvKeyspace, nil
}

// GetSrvKeyspaceNames is part of the topo.Server interface
func (zkts *Server) GetSrvKeyspaceNames(ctx context.Context, cell string) ([]string, error) {
	children, _, err := zkts.zconn.Children(zkPathForCell(cell))
	if err != nil {
		if zookeeper.IsError(err, zookeeper.ZNONODE) {
			return nil, nil
		}
		return nil, err
	}

	sort.Strings(children)
	return children, nil
}

var errSkipUpdate = fmt.Errorf("skip update")

func (zkts *Server) updateTabletEndpoint(oldValue string, oldStat zk.Stat, addr *topo.EndPoint) (newValue string, err error) {
	if oldStat == nil {
		// The incoming object doesn't exist - we haven't been placed in the serving
		// graph yet, so don't update. Assume the next process that rebuilds the graph
		// will get the updated tablet location.
		return "", errSkipUpdate
	}

	var addrs *topo.EndPoints
	if oldValue != "" {
		addrs = &topo.EndPoints{}
		if len(oldValue) > 0 {
			if err := json.Unmarshal([]byte(oldValue), addrs); err != nil {
				return "", fmt.Errorf("EndPoints unmarshal failed: %v %v", oldValue, err)
			}
		}

		foundTablet := false
		for i, entry := range addrs.Entries {
			if entry.Uid == addr.Uid {
				foundTablet = true
				if !topo.EndPointEquality(&entry, addr) {
					addrs.Entries[i] = *addr
				}
				break
			}
		}

		if !foundTablet {
			addrs.Entries = append(addrs.Entries, *addr)
		}
	} else {
		addrs = topo.NewEndPoints()
		addrs.Entries = append(addrs.Entries, *addr)
	}
	return jscfg.ToJSON(addrs), nil
}

// WatchEndPoints is part of the topo.Server interface
func (zkts *Server) WatchEndPoints(ctx context.Context, cell, keyspace, shard string, tabletType topo.TabletType) (<-chan *topo.EndPoints, chan<- struct{}, error) {
	filePath := zkPathForVtName(cell, keyspace, shard, tabletType)

	notifications := make(chan *topo.EndPoints, 10)
	stopWatching := make(chan struct{})

	// waitOrInterrupted will return true if stopWatching is triggered
	waitOrInterrupted := func() bool {
		timer := time.After(WatchSleepDuration)
		select {
		case <-stopWatching:
			close(notifications)
			return true
		case <-timer:
		}
		return false
	}

	go func() {
		for {
			// set the watch
			data, _, watch, err := zkts.zconn.GetW(filePath)
			if err != nil {
				if zookeeper.IsError(err, zookeeper.ZNONODE) {
					// the parent directory doesn't exist
					notifications <- nil
				}

				log.Errorf("Cannot set watch on %v, waiting for %v to retry: %v", filePath, WatchSleepDuration, err)
				if waitOrInterrupted() {
					return
				}
				continue
			}

			// get the initial value, send it, or send nil if no
			// data
			var ep *topo.EndPoints
			sendIt := true
			if len(data) > 0 {
				ep = &topo.EndPoints{}
				if err := json.Unmarshal([]byte(data), ep); err != nil {
					log.Errorf("EndPoints unmarshal failed: %v %v", data, err)
					sendIt = false
				}
			}
			if sendIt {
				notifications <- ep
			}

			// now act on the watch
			select {
			case event, ok := <-watch:
				if !ok {
					log.Warningf("watch on %v was closed, waiting for %v to retry", filePath, WatchSleepDuration)
					if waitOrInterrupted() {
						return
					}
					continue
				}

				if !event.Ok() {
					log.Warningf("received a non-OK event for %v, waiting for %v to retry", filePath, WatchSleepDuration)
					if waitOrInterrupted() {
						return
					}
				}
			case <-stopWatching:
				// user is not interested any more
				close(notifications)
				return
			}
		}
	}()

	return notifications, stopWatching, nil

}
<|endoftext|>