<commit_before>// Copyright 2016 Hajime Hoshi
//
// 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 driver

/*

#cgo LDFLAGS: -llog

#include <android/log.h>
#include <jni.h>
#include <stdlib.h>

// __android_log_print(ANDROID_LOG_ERROR, "NativeCode", "foo", "bar");

static char* initAudioTrack(uintptr_t java_vm, uintptr_t jni_env, jobject context,
    int sampleRate, int channelNum, int bytesPerSample, jobject* audioTrack, int* bufferSize) {
  *bufferSize = 0;
  JavaVM* vm = (JavaVM*)java_vm;
  JNIEnv* env = (JNIEnv*)jni_env;

  const jclass android_media_AudioFormat =
      (*env)->FindClass(env, "android/media/AudioFormat");
  const jclass android_media_AudioManager =
      (*env)->FindClass(env, "android/media/AudioManager");
  const jclass android_media_AudioTrack =
      (*env)->FindClass(env, "android/media/AudioTrack");

  const jint android_media_AudioManager_STREAM_MUSIC =
      (*env)->GetStaticIntField(
          env, android_media_AudioManager,
          (*env)->GetStaticFieldID(env, android_media_AudioManager, "STREAM_MUSIC", "I"));
  const jint android_media_AudioTrack_MODE_STREAM =
      (*env)->GetStaticIntField(
          env, android_media_AudioTrack,
          (*env)->GetStaticFieldID(env, android_media_AudioTrack, "MODE_STREAM", "I"));
  const jint android_media_AudioTrack_WRITE_BLOCKING =
      (*env)->GetStaticIntField(
          env, android_media_AudioTrack,
          (*env)->GetStaticFieldID(env, android_media_AudioTrack, "WRITE_BLOCKING", "I"));
  const jint android_media_AudioFormat_CHANNEL_OUT_MONO =
      (*env)->GetStaticIntField(
          env, android_media_AudioFormat,
          (*env)->GetStaticFieldID(env, android_media_AudioFormat, "CHANNEL_OUT_MONO", "I"));
  const jint android_media_AudioFormat_CHANNEL_OUT_STEREO =
      (*env)->GetStaticIntField(
          env, android_media_AudioFormat,
          (*env)->GetStaticFieldID(env, android_media_AudioFormat, "CHANNEL_OUT_STEREO", "I"));
  const jint android_media_AudioFormat_ENCODING_PCM_8BIT =
      (*env)->GetStaticIntField(
          env, android_media_AudioFormat,
          (*env)->GetStaticFieldID(env, android_media_AudioFormat, "ENCODING_PCM_8BIT", "I"));
  const jint android_media_AudioFormat_ENCODING_PCM_16BIT =
      (*env)->GetStaticIntField(
          env, android_media_AudioFormat,
          (*env)->GetStaticFieldID(env, android_media_AudioFormat, "ENCODING_PCM_16BIT", "I"));

  jint channel = android_media_AudioFormat_CHANNEL_OUT_MONO;
  switch (channelNum) {
  case 1:
    channel = android_media_AudioFormat_CHANNEL_OUT_MONO;
    break;
  case 2:
    channel = android_media_AudioFormat_CHANNEL_OUT_STEREO;
    break;
  default:
    return "invalid channel";
  }

  jint encoding = android_media_AudioFormat_ENCODING_PCM_8BIT;
  switch (bytesPerSample) {
  case 1:
    encoding = android_media_AudioFormat_ENCODING_PCM_8BIT;
    break;
  case 2:
    encoding = android_media_AudioFormat_ENCODING_PCM_16BIT;
    break;
  default:
    return "invalid bytesPerSample";
  }

  *bufferSize =
      (*env)->CallStaticIntMethod(
          env, android_media_AudioTrack,
          (*env)->GetStaticMethodID(env, android_media_AudioTrack, "getMinBufferSize", "(III)I"),
          sampleRate, channel, encoding);

  const jobject tmpAudioTrack =
      (*env)->NewObject(
          env, android_media_AudioTrack,
          (*env)->GetMethodID(env, android_media_AudioTrack, "<init>", "(IIIIII)V"),
          android_media_AudioManager_STREAM_MUSIC,
          sampleRate, channel, encoding, *bufferSize,
          android_media_AudioTrack_MODE_STREAM);
  // Note that *audioTrack will never be released.
  *audioTrack = (*env)->NewGlobalRef(env, tmpAudioTrack);

  // Enqueue empty bytes before playing to avoid underrunning.
  // TODO: Is this really needed? At least, SDL doesn't do the same thing.
  jint writtenBytes = 0;
  do {
    const int length = 1024;
    jbyteArray arr = (*env)->NewByteArray(env, length);
    writtenBytes =
        (*env)->CallIntMethod(
            env, *audioTrack,
            (*env)->GetMethodID(env, android_media_AudioTrack, "write", "([BIII)I"),
            arr, 0, length, android_media_AudioTrack_WRITE_BLOCKING);
  } while (writtenBytes != 0);
  // TODO: Check if writtenBytes < 0

  (*env)->CallVoidMethod(
      env, *audioTrack,
      (*env)->GetMethodID(env, android_media_AudioTrack, "play", "()V"));

  return NULL;
}

static char* writeToAudioTrack(uintptr_t java_vm, uintptr_t jni_env, jobject context,
    jobject audioTrack, int bytesPerSample, void* data, int length) {
  JavaVM* vm = (JavaVM*)java_vm;
  JNIEnv* env = (JNIEnv*)jni_env;

  const jclass android_media_AudioTrack =
      (*env)->FindClass(env, "android/media/AudioTrack");
  const jint android_media_AudioTrack_WRITE_NON_BLOCKING =
      (*env)->GetStaticIntField(
          env, android_media_AudioTrack,
          (*env)->GetStaticFieldID(env, android_media_AudioTrack, "WRITE_NON_BLOCKING", "I"));

  jbyteArray arrInBytes;
  jshortArray arrInShorts;
  switch (bytesPerSample) {
  case 1:
    arrInBytes = (*env)->NewByteArray(env, length);
    (*env)->SetByteArrayRegion(env, arrInBytes, 0, length, data);
    break;
  case 2:
    arrInShorts = (*env)->NewShortArray(env, length);
    (*env)->SetShortArrayRegion(env, arrInShorts, 0, length, data);
    break;
  }
  int i = 0;
  for (i = 0; i < length;) {
    jint result = 0;
    switch (bytesPerSample) {
    case 1:
      result =
          (*env)->CallIntMethod(
              env, audioTrack,
              (*env)->GetMethodID(env, android_media_AudioTrack, "write", "([BIII)I"),
              arrInBytes, i, length - i, android_media_AudioTrack_WRITE_NON_BLOCKING);
      break;
    case 2:
      result =
          (*env)->CallIntMethod(
              env, audioTrack,
              (*env)->GetMethodID(env, android_media_AudioTrack, "write", "([SIII)I"),
              arrInShorts, i, length - i, android_media_AudioTrack_WRITE_NON_BLOCKING);
      break;
    }
    i += result;
  }

  // TODO: Check the result.
  return NULL;
}

*/
import "C"

import (
	"errors"
	"unsafe"
)

type Player struct {
	sampleRate     int
	channelNum     int
	bytesPerSample int
	audioTrack     C.jobject
	buffer         []byte
	bufferSize     int
	chErr          chan error
	chBuffer       chan []byte
}

func NewPlayer(sampleRate, channelNum, bytesPerSample int) (*Player, error) {
	p := &Player{
		sampleRate:     sampleRate,
		channelNum:     channelNum,
		bytesPerSample: bytesPerSample,
		buffer:         []byte{},
		chErr:          make(chan error),
		chBuffer:       make(chan []byte),
	}
	if err := runOnJVM(func(vm, env, ctx uintptr) error {
		audioTrack := C.jobject(nil)
		bufferSize := C.int(0)
		if msg := C.initAudioTrack(C.uintptr_t(vm), C.uintptr_t(env), C.jobject(ctx),
			C.int(sampleRate), C.int(channelNum), C.int(bytesPerSample),
			&audioTrack, &bufferSize); msg != nil {
			return errors.New(C.GoString(msg))
		}
		p.audioTrack = audioTrack
		p.bufferSize = int(bufferSize)
		return nil
	}); err != nil {
		return nil, err
	}
	go func() {
		for bufInBytes := range p.chBuffer {
			var bufInShorts []int16
			if p.bytesPerSample == 2 {
				bufInShorts = make([]int16, len(bufInBytes)/2)
				for i := 0; i < len(bufInShorts); i++ {
					bufInShorts[i] = int16(bufInBytes[2*i]) | (int16(bufInBytes[2*i+1]) << 8)
				}
			}

			if err := runOnJVM(func(vm, env, ctx uintptr) error {
				msg := (*C.char)(nil)
				switch p.bytesPerSample {
				case 1:
					msg = C.writeToAudioTrack(C.uintptr_t(vm), C.uintptr_t(env), C.jobject(ctx),
						p.audioTrack, C.int(p.bytesPerSample),
						unsafe.Pointer(&bufInBytes[0]), C.int(len(bufInBytes)))
				case 2:
					msg = C.writeToAudioTrack(C.uintptr_t(vm), C.uintptr_t(env), C.jobject(ctx),
						p.audioTrack, C.int(p.bytesPerSample),
						unsafe.Pointer(&bufInShorts[0]), C.int(len(bufInShorts)))
				}
				if msg != nil {
					return errors.New(C.GoString(msg))
				}
				return nil
			}); err != nil {
				p.chErr <- err
				return
			}
		}
	}()
	return p, nil
}

func (p *Player) Proceed(data []byte) error {
	p.buffer = append(p.buffer, data...)
	if len(p.buffer) < p.bufferSize {
		return nil
	}
	buf := p.buffer[:p.bufferSize]
	select {
	case p.chBuffer <- buf:
	case err := <-p.chErr:
		return err
	}
	p.buffer = p.buffer[p.bufferSize:]
	return nil
}

func (p *Player) Close() error {
	return nil
}
<commit_msg>audio: Refactoring: Add loop<commit_after>// Copyright 2016 Hajime Hoshi
//
// 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 driver

/*

#cgo LDFLAGS: -llog

#include <android/log.h>
#include <jni.h>
#include <stdlib.h>

// __android_log_print(ANDROID_LOG_ERROR, "NativeCode", "foo", "bar");

static char* initAudioTrack(uintptr_t java_vm, uintptr_t jni_env, jobject context,
    int sampleRate, int channelNum, int bytesPerSample, jobject* audioTrack, int* bufferSize) {
  *bufferSize = 0;
  JavaVM* vm = (JavaVM*)java_vm;
  JNIEnv* env = (JNIEnv*)jni_env;

  const jclass android_media_AudioFormat =
      (*env)->FindClass(env, "android/media/AudioFormat");
  const jclass android_media_AudioManager =
      (*env)->FindClass(env, "android/media/AudioManager");
  const jclass android_media_AudioTrack =
      (*env)->FindClass(env, "android/media/AudioTrack");

  const jint android_media_AudioManager_STREAM_MUSIC =
      (*env)->GetStaticIntField(
          env, android_media_AudioManager,
          (*env)->GetStaticFieldID(env, android_media_AudioManager, "STREAM_MUSIC", "I"));
  const jint android_media_AudioTrack_MODE_STREAM =
      (*env)->GetStaticIntField(
          env, android_media_AudioTrack,
          (*env)->GetStaticFieldID(env, android_media_AudioTrack, "MODE_STREAM", "I"));
  const jint android_media_AudioTrack_WRITE_BLOCKING =
      (*env)->GetStaticIntField(
          env, android_media_AudioTrack,
          (*env)->GetStaticFieldID(env, android_media_AudioTrack, "WRITE_BLOCKING", "I"));
  const jint android_media_AudioFormat_CHANNEL_OUT_MONO =
      (*env)->GetStaticIntField(
          env, android_media_AudioFormat,
          (*env)->GetStaticFieldID(env, android_media_AudioFormat, "CHANNEL_OUT_MONO", "I"));
  const jint android_media_AudioFormat_CHANNEL_OUT_STEREO =
      (*env)->GetStaticIntField(
          env, android_media_AudioFormat,
          (*env)->GetStaticFieldID(env, android_media_AudioFormat, "CHANNEL_OUT_STEREO", "I"));
  const jint android_media_AudioFormat_ENCODING_PCM_8BIT =
      (*env)->GetStaticIntField(
          env, android_media_AudioFormat,
          (*env)->GetStaticFieldID(env, android_media_AudioFormat, "ENCODING_PCM_8BIT", "I"));
  const jint android_media_AudioFormat_ENCODING_PCM_16BIT =
      (*env)->GetStaticIntField(
          env, android_media_AudioFormat,
          (*env)->GetStaticFieldID(env, android_media_AudioFormat, "ENCODING_PCM_16BIT", "I"));

  jint channel = android_media_AudioFormat_CHANNEL_OUT_MONO;
  switch (channelNum) {
  case 1:
    channel = android_media_AudioFormat_CHANNEL_OUT_MONO;
    break;
  case 2:
    channel = android_media_AudioFormat_CHANNEL_OUT_STEREO;
    break;
  default:
    return "invalid channel";
  }

  jint encoding = android_media_AudioFormat_ENCODING_PCM_8BIT;
  switch (bytesPerSample) {
  case 1:
    encoding = android_media_AudioFormat_ENCODING_PCM_8BIT;
    break;
  case 2:
    encoding = android_media_AudioFormat_ENCODING_PCM_16BIT;
    break;
  default:
    return "invalid bytesPerSample";
  }

  *bufferSize =
      (*env)->CallStaticIntMethod(
          env, android_media_AudioTrack,
          (*env)->GetStaticMethodID(env, android_media_AudioTrack, "getMinBufferSize", "(III)I"),
          sampleRate, channel, encoding);

  const jobject tmpAudioTrack =
      (*env)->NewObject(
          env, android_media_AudioTrack,
          (*env)->GetMethodID(env, android_media_AudioTrack, "<init>", "(IIIIII)V"),
          android_media_AudioManager_STREAM_MUSIC,
          sampleRate, channel, encoding, *bufferSize,
          android_media_AudioTrack_MODE_STREAM);
  // Note that *audioTrack will never be released.
  *audioTrack = (*env)->NewGlobalRef(env, tmpAudioTrack);

  // Enqueue empty bytes before playing to avoid underrunning.
  // TODO: Is this really needed? At least, SDL doesn't do the same thing.
  jint writtenBytes = 0;
  do {
    const int length = 1024;
    jbyteArray arr = (*env)->NewByteArray(env, length);
    writtenBytes =
        (*env)->CallIntMethod(
            env, *audioTrack,
            (*env)->GetMethodID(env, android_media_AudioTrack, "write", "([BIII)I"),
            arr, 0, length, android_media_AudioTrack_WRITE_BLOCKING);
  } while (writtenBytes != 0);
  // TODO: Check if writtenBytes < 0

  (*env)->CallVoidMethod(
      env, *audioTrack,
      (*env)->GetMethodID(env, android_media_AudioTrack, "play", "()V"));

  return NULL;
}

static char* writeToAudioTrack(uintptr_t java_vm, uintptr_t jni_env, jobject context,
    jobject audioTrack, int bytesPerSample, void* data, int length) {
  JavaVM* vm = (JavaVM*)java_vm;
  JNIEnv* env = (JNIEnv*)jni_env;

  const jclass android_media_AudioTrack =
      (*env)->FindClass(env, "android/media/AudioTrack");
  const jint android_media_AudioTrack_WRITE_NON_BLOCKING =
      (*env)->GetStaticIntField(
          env, android_media_AudioTrack,
          (*env)->GetStaticFieldID(env, android_media_AudioTrack, "WRITE_NON_BLOCKING", "I"));

  jbyteArray arrInBytes;
  jshortArray arrInShorts;
  switch (bytesPerSample) {
  case 1:
    arrInBytes = (*env)->NewByteArray(env, length);
    (*env)->SetByteArrayRegion(env, arrInBytes, 0, length, data);
    break;
  case 2:
    arrInShorts = (*env)->NewShortArray(env, length);
    (*env)->SetShortArrayRegion(env, arrInShorts, 0, length, data);
    break;
  }
  int i = 0;
  for (i = 0; i < length;) {
    jint result = 0;
    switch (bytesPerSample) {
    case 1:
      result =
          (*env)->CallIntMethod(
              env, audioTrack,
              (*env)->GetMethodID(env, android_media_AudioTrack, "write", "([BIII)I"),
              arrInBytes, i, length - i, android_media_AudioTrack_WRITE_NON_BLOCKING);
      break;
    case 2:
      result =
          (*env)->CallIntMethod(
              env, audioTrack,
              (*env)->GetMethodID(env, android_media_AudioTrack, "write", "([SIII)I"),
              arrInShorts, i, length - i, android_media_AudioTrack_WRITE_NON_BLOCKING);
      break;
    }
    i += result;
  }

  // TODO: Check the result.
  return NULL;
}

*/
import "C"

import (
	"errors"
	"unsafe"
)

type Player struct {
	sampleRate     int
	channelNum     int
	bytesPerSample int
	audioTrack     C.jobject
	buffer         []byte
	bufferSize     int
	chErr          chan error
	chBuffer       chan []byte
}

func NewPlayer(sampleRate, channelNum, bytesPerSample int) (*Player, error) {
	p := &Player{
		sampleRate:     sampleRate,
		channelNum:     channelNum,
		bytesPerSample: bytesPerSample,
		buffer:         []byte{},
		chErr:          make(chan error),
		chBuffer:       make(chan []byte),
	}
	if err := runOnJVM(func(vm, env, ctx uintptr) error {
		audioTrack := C.jobject(nil)
		bufferSize := C.int(0)
		if msg := C.initAudioTrack(C.uintptr_t(vm), C.uintptr_t(env), C.jobject(ctx),
			C.int(sampleRate), C.int(channelNum), C.int(bytesPerSample),
			&audioTrack, &bufferSize); msg != nil {
			return errors.New(C.GoString(msg))
		}
		p.audioTrack = audioTrack
		p.bufferSize = int(bufferSize)
		return nil
	}); err != nil {
		return nil, err
	}
	go p.loop()
	return p, nil
}

func (p *Player) loop() {
	for bufInBytes := range p.chBuffer {
		var bufInShorts []int16
		if p.bytesPerSample == 2 {
			bufInShorts = make([]int16, len(bufInBytes)/2)
			for i := 0; i < len(bufInShorts); i++ {
				bufInShorts[i] = int16(bufInBytes[2*i]) | (int16(bufInBytes[2*i+1]) << 8)
			}
		}

		if err := runOnJVM(func(vm, env, ctx uintptr) error {
			msg := (*C.char)(nil)
			switch p.bytesPerSample {
			case 1:
				msg = C.writeToAudioTrack(C.uintptr_t(vm), C.uintptr_t(env), C.jobject(ctx),
					p.audioTrack, C.int(p.bytesPerSample),
					unsafe.Pointer(&bufInBytes[0]), C.int(len(bufInBytes)))
			case 2:
				msg = C.writeToAudioTrack(C.uintptr_t(vm), C.uintptr_t(env), C.jobject(ctx),
					p.audioTrack, C.int(p.bytesPerSample),
					unsafe.Pointer(&bufInShorts[0]), C.int(len(bufInShorts)))
			}
			if msg != nil {
				return errors.New(C.GoString(msg))
			}
			return nil
		}); err != nil {
			p.chErr <- err
			return
		}
	}
}

func (p *Player) Proceed(data []byte) error {
	p.buffer = append(p.buffer, data...)
	if len(p.buffer) < p.bufferSize {
		return nil
	}
	buf := p.buffer[:p.bufferSize]
	select {
	case p.chBuffer <- buf:
	case err := <-p.chErr:
		return err
	}
	p.buffer = p.buffer[p.bufferSize:]
	return nil
}

func (p *Player) Close() error {
	return nil
}
<|endoftext|>


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

// HTTP reverse proxy handler

package proxy

import (
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

// onExitFlushLoop is a callback set by tests to detect the state of the
// flushLoop() goroutine.
var onExitFlushLoop func()

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// Director must be a function which modifies
	// the request into a new request to be sent
	// using Transport. Its response is then copied
	// back to the original client unmodified.
	Director func(*http.Request)

	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration

	// ErrorLog specifies an optional logger for errors
	// that occur when attempting to proxy the request.
	// If nil, logging goes to os.Stderr via the log package's
	// standard logger.
	ErrorLog *log.Logger
}

func singleJoiningSlash(a, b string) string {
	aslash := strings.HasSuffix(a, "/")
	bslash := strings.HasPrefix(b, "/")
	switch {
	case aslash && bslash:
		return a + b[1:]
	case !aslash && !bslash:
		return a + "/" + b
	}
	return a + b
}

// NewSingleHostReverseProxy returns a new ReverseProxy that rewrites
// URLs to the scheme, host, and base path provided in target. If the
// target's path is "/base" and the incoming request was for "/dir",
// the target request will be for /base/dir.
func NewSingleHostReverseProxy(target *url.URL) *ReverseProxy {
	targetQuery := target.RawQuery
	director := func(req *http.Request) {
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		req.URL.Path = singleJoiningSlash(target.Path, req.URL.Path)
		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}
	}
	return &ReverseProxy{Director: director}
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	transport := p.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}

	outreq := new(http.Request)
	*outreq = *req // includes shallow copies of maps, but okay

	p.Director(outreq)
	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	// Remove hop-by-hop headers to the backend.  Especially
	// important is "Connection" because we want a persistent
	// connection, regardless of what the client sent to us.  This
	// is modifying the same underlying map from req (shallow
	// copied above) so we only copy it if necessary.
	copiedHeaders := false
	for _, h := range hopHeaders {
		if outreq.Header.Get(h) != "" {
			if !copiedHeaders {
				outreq.Header = make(http.Header)
				copyHeader(outreq.Header, req.Header)
				copiedHeaders = true
			}
			outreq.Header.Del(h)
		}
	}

	if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
		// If we aren't the first proxy retain prior
		// X-Forwarded-For information as a comma+space
		// separated list and fold multiple headers into one.
		if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		outreq.Header.Set("X-Forwarded-For", clientIP)
	}

	res, err := transport.RoundTrip(outreq)
	if err != nil {
		p.logf("http: proxy error: %v", err)
		rw.WriteHeader(http.StatusInternalServerError)
		return
	}
	defer res.Body.Close()

	for _, h := range hopHeaders {
		res.Header.Del(h)
	}

	copyHeader(rw.Header(), res.Header)

	rw.WriteHeader(res.StatusCode)
	p.copyResponse(rw, res.Body)
}

func (p *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	if p.FlushInterval != 0 {
		if wf, ok := dst.(writeFlusher); ok {
			mlw := &maxLatencyWriter{
				dst:     wf,
				latency: p.FlushInterval,
				done:    make(chan bool),
			}
			go mlw.flushLoop()
			defer mlw.stop()
			dst = mlw
		}
	}

	io.Copy(dst, src)
}

func (p *ReverseProxy) logf(format string, args ...interface{}) {
	if p.ErrorLog != nil {
		p.ErrorLog.Printf(format, args...)
	} else {
		log.Printf(format, args...)
	}
}

type writeFlusher interface {
	io.Writer
	http.Flusher
}

type maxLatencyWriter struct {
	dst     writeFlusher
	latency time.Duration

	lk   sync.Mutex // protects Write + Flush
	done chan bool
}

func (m *maxLatencyWriter) Write(p []byte) (int, error) {
	m.lk.Lock()
	defer m.lk.Unlock()
	return m.dst.Write(p)
}

func (m *maxLatencyWriter) flushLoop() {
	t := time.NewTicker(m.latency)
	defer t.Stop()
	for {
		select {
		case <-m.done:
			if onExitFlushLoop != nil {
				onExitFlushLoop()
			}
			return
		case <-t.C:
			m.lk.Lock()
			m.dst.Flush()
			m.lk.Unlock()
		}
	}
}

func (m *maxLatencyWriter) stop() { m.done <- true }
<commit_msg>router: remove unused parts of ReverseProxy from stdlib<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// HTTP reverse proxy handler

package proxy

import (
	"io"
	"log"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
)

// onExitFlushLoop is a callback set by tests to detect the state of the
// flushLoop() goroutine.
var onExitFlushLoop func()

// ReverseProxy is an HTTP Handler that takes an incoming request and
// sends it to another server, proxying the response back to the
// client.
type ReverseProxy struct {
	// The transport used to perform proxy requests.
	// If nil, http.DefaultTransport is used.
	Transport http.RoundTripper

	// FlushInterval specifies the flush interval
	// to flush to the client while copying the
	// response body.
	// If zero, no periodic flushing is done.
	FlushInterval time.Duration

	// ErrorLog specifies an optional logger for errors
	// that occur when attempting to proxy the request.
	// If nil, logging goes to os.Stderr via the log package's
	// standard logger.
	ErrorLog *log.Logger
}

func copyHeader(dst, src http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	transport := p.Transport
	if transport == nil {
		transport = http.DefaultTransport
	}

	outreq := new(http.Request)
	*outreq = *req // includes shallow copies of maps, but okay

	outreq.Proto = "HTTP/1.1"
	outreq.ProtoMajor = 1
	outreq.ProtoMinor = 1
	outreq.Close = false

	// Remove hop-by-hop headers to the backend.  Especially
	// important is "Connection" because we want a persistent
	// connection, regardless of what the client sent to us.  This
	// is modifying the same underlying map from req (shallow
	// copied above) so we only copy it if necessary.
	copiedHeaders := false
	for _, h := range hopHeaders {
		if outreq.Header.Get(h) != "" {
			if !copiedHeaders {
				outreq.Header = make(http.Header)
				copyHeader(outreq.Header, req.Header)
				copiedHeaders = true
			}
			outreq.Header.Del(h)
		}
	}

	if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
		// If we aren't the first proxy retain prior
		// X-Forwarded-For information as a comma+space
		// separated list and fold multiple headers into one.
		if prior, ok := outreq.Header["X-Forwarded-For"]; ok {
			clientIP = strings.Join(prior, ", ") + ", " + clientIP
		}
		outreq.Header.Set("X-Forwarded-For", clientIP)
	}

	res, err := transport.RoundTrip(outreq)
	if err != nil {
		p.logf("http: proxy error: %v", err)
		rw.WriteHeader(http.StatusInternalServerError)
		return
	}
	defer res.Body.Close()

	for _, h := range hopHeaders {
		res.Header.Del(h)
	}

	copyHeader(rw.Header(), res.Header)

	rw.WriteHeader(res.StatusCode)
	p.copyResponse(rw, res.Body)
}

func (p *ReverseProxy) copyResponse(dst io.Writer, src io.Reader) {
	if p.FlushInterval != 0 {
		if wf, ok := dst.(writeFlusher); ok {
			mlw := &maxLatencyWriter{
				dst:     wf,
				latency: p.FlushInterval,
				done:    make(chan bool),
			}
			go mlw.flushLoop()
			defer mlw.stop()
			dst = mlw
		}
	}

	io.Copy(dst, src)
}

func (p *ReverseProxy) logf(format string, args ...interface{}) {
	if p.ErrorLog != nil {
		p.ErrorLog.Printf(format, args...)
	} else {
		log.Printf(format, args...)
	}
}

type writeFlusher interface {
	io.Writer
	http.Flusher
}

type maxLatencyWriter struct {
	dst     writeFlusher
	latency time.Duration

	lk   sync.Mutex // protects Write + Flush
	done chan bool
}

func (m *maxLatencyWriter) Write(p []byte) (int, error) {
	m.lk.Lock()
	defer m.lk.Unlock()
	return m.dst.Write(p)
}

func (m *maxLatencyWriter) flushLoop() {
	t := time.NewTicker(m.latency)
	defer t.Stop()
	for {
		select {
		case <-m.done:
			if onExitFlushLoop != nil {
				onExitFlushLoop()
			}
			return
		case <-t.C:
			m.lk.Lock()
			m.dst.Flush()
			m.lk.Unlock()
		}
	}
}

func (m *maxLatencyWriter) stop() { m.done <- true }
<|endoftext|>


<commit_before>package router

import (
	"net/http"
	"strconv"

	"github.com/ewhal/nyaa/db"
	"github.com/ewhal/nyaa/model"
	"github.com/ewhal/nyaa/service/captcha"
	"github.com/ewhal/nyaa/service/torrent"
	"github.com/ewhal/nyaa/util/log"
	"github.com/gorilla/mux"
)

func ViewHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]

	torrent, err := torrentService.GetTorrentById(id)
	if err != nil {
		NotFoundHandler(w, r)
		return
	}
	b := torrent.ToJson()
	htv := ViewTemplateVariables{b, captcha.Captcha{CaptchaID: captcha.GetID()}, NewSearchForm(), Navigation{}, GetUser(r), r.URL, mux.CurrentRoute(r)}

	err = viewTemplate.ExecuteTemplate(w, "index.html", htv)
	if err != nil {
		log.Errorf("ViewHandler(): %s", err)
	}
}

func PostCommentHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]

	userCaptcha := captcha.Extract(r)
	if !captcha.Authenticate(userCaptcha) {
		http.Error(w, "bad captcha", 403)
	}
	content := r.FormValue("comment")

	idNum, err := strconv.Atoi(id)
	comment := model.Comment{Username: "れんちょん", Content: content, TorrentId: idNum}
	db.ORM.Create(&comment)

	url, err := Router.Get("view_torrent").URL("id", id)
	if err == nil {
		http.Redirect(w, r, url.String(), 302)
	} else {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}
<commit_msg>sanitize comments<commit_after>package router

import (
	"net/http"
	"strconv"

	"github.com/ewhal/nyaa/db"
	"github.com/ewhal/nyaa/model"
	"github.com/ewhal/nyaa/service/captcha"
	"github.com/ewhal/nyaa/service/torrent"
	"github.com/ewhal/nyaa/util/log"
	"github.com/gorilla/mux"
    "github.com/microcosm-cc/bluemonday"
)

var p = bluemonday.UGCPolicy()

func ViewHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]

	torrent, err := torrentService.GetTorrentById(id)
	if err != nil {
		NotFoundHandler(w, r)
		return
	}
	b := torrent.ToJson()
	htv := ViewTemplateVariables{b, captcha.Captcha{CaptchaID: captcha.GetID()}, NewSearchForm(), Navigation{}, GetUser(r), r.URL, mux.CurrentRoute(r)}

	err = viewTemplate.ExecuteTemplate(w, "index.html", htv)
	if err != nil {
		log.Errorf("ViewHandler(): %s", err)
	}
}

func PostCommentHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["id"]

	userCaptcha := captcha.Extract(r)
	if !captcha.Authenticate(userCaptcha) {
		http.Error(w, "bad captcha", 403)
	}
	content := p.Sanitize(r.FormValue("comment"))

	idNum, err := strconv.Atoi(id)
	comment := model.Comment{Username: "れんちょん", Content: content, TorrentId: idNum}
	db.ORM.Create(&comment)

	url, err := Router.Get("view_torrent").URL("id", id)
	if err == nil {
		http.Redirect(w, r, url.String(), 302)
	} else {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Stratumn SAS. All rights reserved.
// Use of this source code is governed by the license
// that can be found in the LICENSE file.

// Package batchfossilizer implements a fossilizer that fossilize batches of data using a Merkle tree.
// The evidence will contain the Merkle root, the Merkle path, and a timestamp.
package batchfossilizer

import (
	"encoding/gob"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/stratumn/go/fossilizer"
	"github.com/stratumn/go/types"
	"github.com/stratumn/goprivate/merkle"
)

const (
	// Name is the name set in the fossilizer's information.
	Name = "batch"

	// Description is the description set in the fossilizer's information.
	Description = "Stratumn Batch Fossilizer"

	// DefaultInterval is the default interval between batches.
	DefaultInterval = 10 * time.Minute

	// DefaultMaxLeaves if the default maximum number of leaves of a Merkle tree.
	DefaultMaxLeaves = 32 * 1024

	// DefaultMaxSimBatches is the default maximum number of simultaneous batches.
	DefaultMaxSimBatches = 1

	// DefaultArchive is whether to archive completed batches by default.
	DefaultArchive = true

	// DefaultStopBatch is whether to do a batch on stop by default.
	DefaultStopBatch = true

	// DefaultFSync is whether to fsync after saving a hash to disk by default.
	DefaultFSync = false

	// PendingExt is the pending hashes filename extension.
	PendingExt = "pending"

	// DirPerm is the directory's permissions.
	DirPerm = 0644

	// FilePerm is the files's permissions.
	FilePerm = 0644
)

// Config contains configuration options for the fossilizer.
type Config struct {
	// A version string that will be set in the store's information.
	Version string

	// A git commit sha that will be set in the store's information.
	Commit string

	// Interval between batches.
	Interval time.Duration

	// Maximum number of leaves of a Merkle tree.
	MaxLeaves int

	// Maximum number of simultaneous batches.
	MaxSimBatches int

	// Where to store pending hashes.
	// If empty, pending hashes are not saved and will be lost if stopped abruptly.
	Path string

	// Whether to archive completed batches.
	Archive bool

	// Whether to do a batch on stop.
	StopBatch bool

	// Whether to fsync after saving a hash to disk.
	FSync bool
}

// Info is the info returned by GetInfo.
type Info struct {
	Name        string `json:"name"`
	Description string `json:"description"`
	Version     string `json:"version"`
	Commit      string `json:"commit"`
}

// Evidence is the evidence sent to the result channel.
type Evidence struct {
	Time int64          `json:"time"`
	Root *types.Bytes32 `json:"merkleRoot"`
	Path merkle.Path    `json:"merklePath"`
}

// EvidenceWrapper wraps evidence with a namespace.
type EvidenceWrapper struct {
	Evidence *Evidence `json:"batch"`
}

type batch struct {
	leaves []types.Bytes32
	meta   [][]byte
	path   string
}

type chunk struct {
	Data []byte
	Meta []byte
}

// Fossilizer is the type that implements github.com/stratumn/go/fossilizer.Adapter.
type Fossilizer struct {
	config      *Config
	resultChans []chan *fossilizer.Result
	leaves      []types.Bytes32
	meta        [][]byte
	file        *os.File
	encoder     *gob.Encoder
	mutex       sync.Mutex
	waitGroup   sync.WaitGroup
	sem         chan struct{}
	closeChan   chan error
}

// New creates an instance of a Fossilizer.
func New(config *Config) (*Fossilizer, error) {
	maxLeaves := config.MaxLeaves
	if maxLeaves == 0 {
		maxLeaves = DefaultMaxLeaves
	}

	maxSimBatches := config.MaxSimBatches
	if maxSimBatches == 0 {
		maxSimBatches = DefaultMaxSimBatches
	}

	a := &Fossilizer{
		config:    config,
		leaves:    make([]types.Bytes32, 0, maxLeaves),
		meta:      make([][]byte, 0, maxLeaves),
		sem:       make(chan struct{}, maxSimBatches),
		closeChan: make(chan error),
	}

	if a.config.Path != "" {
		if err := os.MkdirAll(a.config.Path, DirPerm); err != nil {
			return nil, err
		}

		if err := a.recover(); err != nil {
			return nil, err
		}
	}

	return a, nil
}

// GetInfo implements github.com/stratumn/go/fossilizer.Adapter.GetInfo.
func (a *Fossilizer) GetInfo() (interface{}, error) {
	return &Info{
		Name:        Name,
		Description: Description,
		Version:     a.config.Version,
		Commit:      a.config.Commit,
	}, nil
}

// AddResultChan implements github.com/stratumn/go/fossilizer.Adapter.AddResultChan.
func (a *Fossilizer) AddResultChan(resultChan chan *fossilizer.Result) {
	a.resultChans = append(a.resultChans, resultChan)
}

// Fossilize implements github.com/stratumn/go/fossilizer.Adapter.Fossilize.
func (a *Fossilizer) Fossilize(data []byte, meta []byte) error {
	a.mutex.Lock()
	defer a.mutex.Unlock()

	if a.closeChan == nil {
		return errors.New("fossilizer is stopped")
	}

	if a.config.Path != "" {
		if a.file == nil {
			if err := a.open(); err != nil {
				return err
			}
		}

		if err := a.write(data, meta); err != nil {
			return err
		}
	}

	var leaf types.Bytes32
	copy(leaf[:], data)
	a.leaves = append(a.leaves, leaf)
	a.meta = append(a.meta, meta)

	maxLeaves := a.config.MaxLeaves
	if maxLeaves == 0 {
		maxLeaves = DefaultMaxLeaves
	}
	if len(a.leaves) >= maxLeaves {
		if err := a.requestBatch(); err != nil {
			a.closeChan <- err
			return err
		}
		log.Printf("Requested new batch because the maximum number of leaves (%d) was reached", maxLeaves)
	}

	return nil
}

// Start starts the fossilizer.
func (a *Fossilizer) Start() error {
	interval := a.config.Interval
	if interval == 0 {
		interval = DefaultInterval
	}

	for {
		select {
		case <-time.After(interval):
			a.mutex.Lock()
			if len(a.leaves) > 0 {
				if err := a.requestBatch(); err != nil {
					return err
				}
				log.Printf("Requested new batch because the %s interval was reached", interval)
			} else {
				log.Printf("No batch is needed after the %s interval because there are no pending hashes", interval)
			}
			a.mutex.Unlock()
		case err := <-a.closeChan:
			return err
		}
	}
}

// Stop stops the fossilizer.
func (a *Fossilizer) Stop() error {
	a.mutex.Lock()
	defer a.mutex.Unlock()

	close(a.closeChan)
	a.closeChan = nil

	if a.config.StopBatch {
		if len(a.leaves) > 0 {
			if err := a.requestBatch(); err != nil {
				return err
			}
			log.Print("Requested final batch for pending hashes")
		} else {
			log.Print("No final batch is needed because there are no pending hashes")
		}
	}

	a.waitGroup.Wait()
	close(a.sem)

	if a.file != nil {
		return a.file.Close()
	}

	return nil
}

func (a *Fossilizer) batch(b batch) {
	defer func() {
		a.waitGroup.Done()
		<-a.sem
	}()

	a.sem <- struct{}{}

	tree, err := merkle.NewStaticTree(b.leaves)

	if err != nil {
		a.closeChan <- err
		return
	}

	var (
		meta = b.meta
		ts   = time.Now().UTC().Unix()
		root = tree.Root()
	)

	log.Printf("Created tree with Merkle root %q", root)

	for i := 0; i < tree.LeavesLen(); i++ {
		leaf := tree.Leaf(i)
		r := &fossilizer.Result{
			Evidence: &EvidenceWrapper{
				&Evidence{
					Time: ts,
					Root: root,
					Path: tree.Path(i),
				},
			},
			Data: leaf[:],
			Meta: meta[i],
		}

		for _, c := range a.resultChans {
			c <- r
		}
	}

	log.Printf("Sent evidence for batch with Merkle root %q", root)

	if b.path != "" {
		if a.config.Archive {
			path := filepath.Join(a.config.Path, root.String())
			if err := os.Rename(b.path, path); err != nil {
				log.Printf("Error: %s", err)
			}
			log.Printf("Renamed pending hashes file %q to %q", filepath.Base(b.path), filepath.Base(path))
		} else {
			if err := os.Remove(b.path); err != nil {
				log.Printf("Error: %s", err)
			}
			log.Printf("Removed pending hashes file %q", filepath.Base(b.path))
		}
	}

	log.Printf("Finished batch with Merkle root %q", root)
}

func (a *Fossilizer) requestBatch() error {
	var path string

	if a.file != nil {
		path = a.file.Name()
		if err := a.file.Close(); err != nil {
			return err
		}
		a.file = nil
	}

	a.waitGroup.Add(1)
	go a.batch(batch{a.leaves, a.meta, path})

	maxLeaves := a.config.MaxLeaves
	if maxLeaves == 0 {
		maxLeaves = DefaultMaxLeaves
	}

	a.leaves, a.meta = make([]types.Bytes32, 0, maxLeaves), make([][]byte, 0, maxLeaves)

	return nil
}

func (a *Fossilizer) recover() error {
	matches, err := filepath.Glob(filepath.Join(a.config.Path, "*."+PendingExt))
	if err != nil {
		return err
	}

	for _, path := range matches {
		file, err := os.OpenFile(path, os.O_RDONLY|os.O_EXCL, FilePerm)
		if err != nil {
			return err
		}

		dec := gob.NewDecoder(file)

		for {
			var c chunk

			if err := dec.Decode(&c); err != nil {
				if err == io.EOF {
					break
				}
				file.Close()
				return err
			}

			if err := a.Fossilize(c.Data, c.Meta); err != nil {
				file.Close()
				return err
			}
		}

		a.waitGroup.Wait()
		file.Close()

		if err := os.Remove(path); err != nil {
			return err
		}

		log.Printf("Recovered pending hashes file %q", filepath.Base(path))
	}

	return nil
}

func (a *Fossilizer) open() error {
	filename := fmt.Sprintf("%d.%s", time.Now().UTC().UnixNano(), PendingExt)
	path := filepath.Join(a.config.Path, filename)
	file, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY|os.O_EXCL|os.O_CREATE, FilePerm)
	if err != nil {
		return err
	}
	log.Printf("Opened pending hashes file %q", filepath.Base(path))

	a.file, a.encoder = file, gob.NewEncoder(file)

	return nil
}

func (a *Fossilizer) write(data []byte, meta []byte) error {
	if err := a.encoder.Encode(chunk{data, meta}); err != nil {
		return err
	}

	if a.config.FSync {
		if err := a.file.Sync(); err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>batchfossilizer: Fix batch interval<commit_after>// Copyright 2016 Stratumn SAS. All rights reserved.
// Use of this source code is governed by the license
// that can be found in the LICENSE file.

// Package batchfossilizer implements a fossilizer that fossilize batches of data using a Merkle tree.
// The evidence will contain the Merkle root, the Merkle path, and a timestamp.
package batchfossilizer

import (
	"encoding/gob"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/stratumn/go/fossilizer"
	"github.com/stratumn/go/types"
	"github.com/stratumn/goprivate/merkle"
)

const (
	// Name is the name set in the fossilizer's information.
	Name = "batch"

	// Description is the description set in the fossilizer's information.
	Description = "Stratumn Batch Fossilizer"

	// DefaultInterval is the default interval between batches.
	DefaultInterval = 10 * time.Minute

	// DefaultMaxLeaves if the default maximum number of leaves of a Merkle tree.
	DefaultMaxLeaves = 32 * 1024

	// DefaultMaxSimBatches is the default maximum number of simultaneous batches.
	DefaultMaxSimBatches = 1

	// DefaultArchive is whether to archive completed batches by default.
	DefaultArchive = true

	// DefaultStopBatch is whether to do a batch on stop by default.
	DefaultStopBatch = true

	// DefaultFSync is whether to fsync after saving a hash to disk by default.
	DefaultFSync = false

	// PendingExt is the pending hashes filename extension.
	PendingExt = "pending"

	// DirPerm is the directory's permissions.
	DirPerm = 0644

	// FilePerm is the files's permissions.
	FilePerm = 0644
)

// Config contains configuration options for the fossilizer.
type Config struct {
	// A version string that will be set in the store's information.
	Version string

	// A git commit sha that will be set in the store's information.
	Commit string

	// Interval between batches.
	Interval time.Duration

	// Maximum number of leaves of a Merkle tree.
	MaxLeaves int

	// Maximum number of simultaneous batches.
	MaxSimBatches int

	// Where to store pending hashes.
	// If empty, pending hashes are not saved and will be lost if stopped abruptly.
	Path string

	// Whether to archive completed batches.
	Archive bool

	// Whether to do a batch on stop.
	StopBatch bool

	// Whether to fsync after saving a hash to disk.
	FSync bool
}

// Info is the info returned by GetInfo.
type Info struct {
	Name        string `json:"name"`
	Description string `json:"description"`
	Version     string `json:"version"`
	Commit      string `json:"commit"`
}

// Evidence is the evidence sent to the result channel.
type Evidence struct {
	Time int64          `json:"time"`
	Root *types.Bytes32 `json:"merkleRoot"`
	Path merkle.Path    `json:"merklePath"`
}

// EvidenceWrapper wraps evidence with a namespace.
type EvidenceWrapper struct {
	Evidence *Evidence `json:"batch"`
}

type batch struct {
	leaves []types.Bytes32
	meta   [][]byte
	path   string
}

type chunk struct {
	Data []byte
	Meta []byte
}

// Fossilizer is the type that implements github.com/stratumn/go/fossilizer.Adapter.
type Fossilizer struct {
	config      *Config
	resultChans []chan *fossilizer.Result
	leaves      []types.Bytes32
	meta        [][]byte
	file        *os.File
	encoder     *gob.Encoder
	mutex       sync.Mutex
	waitGroup   sync.WaitGroup
	sem         chan struct{}
	resetChan   chan struct{}
	closeChan   chan error
}

// New creates an instance of a Fossilizer.
func New(config *Config) (*Fossilizer, error) {
	maxLeaves := config.MaxLeaves
	if maxLeaves == 0 {
		maxLeaves = DefaultMaxLeaves
	}

	maxSimBatches := config.MaxSimBatches
	if maxSimBatches == 0 {
		maxSimBatches = DefaultMaxSimBatches
	}

	a := &Fossilizer{
		config:    config,
		leaves:    make([]types.Bytes32, 0, maxLeaves),
		meta:      make([][]byte, 0, maxLeaves),
		sem:       make(chan struct{}, maxSimBatches),
		resetChan: make(chan struct{}),
		closeChan: make(chan error),
	}

	if a.config.Path != "" {
		if err := os.MkdirAll(a.config.Path, DirPerm); err != nil {
			return nil, err
		}

		if err := a.recover(); err != nil {
			return nil, err
		}
	}

	return a, nil
}

// GetInfo implements github.com/stratumn/go/fossilizer.Adapter.GetInfo.
func (a *Fossilizer) GetInfo() (interface{}, error) {
	return &Info{
		Name:        Name,
		Description: Description,
		Version:     a.config.Version,
		Commit:      a.config.Commit,
	}, nil
}

// AddResultChan implements github.com/stratumn/go/fossilizer.Adapter.AddResultChan.
func (a *Fossilizer) AddResultChan(resultChan chan *fossilizer.Result) {
	a.resultChans = append(a.resultChans, resultChan)
}

// Fossilize implements github.com/stratumn/go/fossilizer.Adapter.Fossilize.
func (a *Fossilizer) Fossilize(data []byte, meta []byte) error {
	a.mutex.Lock()
	defer a.mutex.Unlock()

	if a.closeChan == nil {
		return errors.New("fossilizer is stopped")
	}

	if a.config.Path != "" {
		if a.file == nil {
			if err := a.open(); err != nil {
				return err
			}
		}

		if err := a.write(data, meta); err != nil {
			return err
		}
	}

	var leaf types.Bytes32
	copy(leaf[:], data)
	a.leaves = append(a.leaves, leaf)
	a.meta = append(a.meta, meta)

	maxLeaves := a.config.MaxLeaves
	if maxLeaves == 0 {
		maxLeaves = DefaultMaxLeaves
	}
	if len(a.leaves) >= maxLeaves {
		if err := a.requestBatch(); err != nil {
			a.closeChan <- err
			return err
		}
		a.resetChan <- struct{}{}
		log.Printf("Requested new batch because the maximum number of leaves (%d) was reached", maxLeaves)
	}

	return nil
}

// Start starts the fossilizer.
func (a *Fossilizer) Start() error {
	interval := a.config.Interval
	if interval == 0 {
		interval = DefaultInterval
	}

	for {
		select {
		case <-time.After(interval):
			a.mutex.Lock()
			if len(a.leaves) > 0 {
				if err := a.requestBatch(); err != nil {
					return err
				}
				log.Printf("Requested new batch because the %s interval was reached", interval)
			} else {
				log.Printf("No batch is needed after the %s interval because there are no pending hashes", interval)
			}
			a.mutex.Unlock()
		case <-a.resetChan:
		case err := <-a.closeChan:
			return err
		}
	}
}

// Stop stops the fossilizer.
func (a *Fossilizer) Stop() error {
	a.mutex.Lock()
	defer a.mutex.Unlock()

	close(a.closeChan)
	a.closeChan = nil

	if a.config.StopBatch {
		if len(a.leaves) > 0 {
			if err := a.requestBatch(); err != nil {
				return err
			}
			log.Print("Requested final batch for pending hashes")
		} else {
			log.Print("No final batch is needed because there are no pending hashes")
		}
	}

	a.waitGroup.Wait()
	close(a.sem)
	close(a.resetChan)

	if a.file != nil {
		return a.file.Close()
	}

	return nil
}

func (a *Fossilizer) batch(b batch) {
	defer func() {
		a.waitGroup.Done()
		<-a.sem
	}()

	a.sem <- struct{}{}

	tree, err := merkle.NewStaticTree(b.leaves)

	if err != nil {
		a.closeChan <- err
		return
	}

	var (
		meta = b.meta
		ts   = time.Now().UTC().Unix()
		root = tree.Root()
	)

	log.Printf("Created tree with Merkle root %q", root)

	for i := 0; i < tree.LeavesLen(); i++ {
		leaf := tree.Leaf(i)
		r := &fossilizer.Result{
			Evidence: &EvidenceWrapper{
				&Evidence{
					Time: ts,
					Root: root,
					Path: tree.Path(i),
				},
			},
			Data: leaf[:],
			Meta: meta[i],
		}

		for _, c := range a.resultChans {
			c <- r
		}
	}

	log.Printf("Sent evidence for batch with Merkle root %q", root)

	if b.path != "" {
		if a.config.Archive {
			path := filepath.Join(a.config.Path, root.String())
			if err := os.Rename(b.path, path); err != nil {
				log.Printf("Error: %s", err)
			}
			log.Printf("Renamed pending hashes file %q to %q", filepath.Base(b.path), filepath.Base(path))
		} else {
			if err := os.Remove(b.path); err != nil {
				log.Printf("Error: %s", err)
			}
			log.Printf("Removed pending hashes file %q", filepath.Base(b.path))
		}
	}

	log.Printf("Finished batch with Merkle root %q", root)
}

func (a *Fossilizer) requestBatch() error {
	var path string

	if a.file != nil {
		path = a.file.Name()
		if err := a.file.Close(); err != nil {
			return err
		}
		a.file = nil
	}

	a.waitGroup.Add(1)
	go a.batch(batch{a.leaves, a.meta, path})

	maxLeaves := a.config.MaxLeaves
	if maxLeaves == 0 {
		maxLeaves = DefaultMaxLeaves
	}

	a.leaves, a.meta = make([]types.Bytes32, 0, maxLeaves), make([][]byte, 0, maxLeaves)

	return nil
}

func (a *Fossilizer) recover() error {
	matches, err := filepath.Glob(filepath.Join(a.config.Path, "*."+PendingExt))
	if err != nil {
		return err
	}

	for _, path := range matches {
		file, err := os.OpenFile(path, os.O_RDONLY|os.O_EXCL, FilePerm)
		if err != nil {
			return err
		}

		dec := gob.NewDecoder(file)

		for {
			var c chunk

			if err := dec.Decode(&c); err != nil {
				if err == io.EOF {
					break
				}
				file.Close()
				return err
			}

			if err := a.Fossilize(c.Data, c.Meta); err != nil {
				file.Close()
				return err
			}
		}

		a.waitGroup.Wait()
		file.Close()

		if err := os.Remove(path); err != nil {
			return err
		}

		log.Printf("Recovered pending hashes file %q", filepath.Base(path))
	}

	return nil
}

func (a *Fossilizer) open() error {
	filename := fmt.Sprintf("%d.%s", time.Now().UTC().UnixNano(), PendingExt)
	path := filepath.Join(a.config.Path, filename)
	file, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY|os.O_EXCL|os.O_CREATE, FilePerm)
	if err != nil {
		return err
	}
	log.Printf("Opened pending hashes file %q", filepath.Base(path))

	a.file, a.encoder = file, gob.NewEncoder(file)

	return nil
}

func (a *Fossilizer) write(data []byte, meta []byte) error {
	if err := a.encoder.Encode(chunk{data, meta}); err != nil {
		return err
	}

	if a.config.FSync {
		if err := a.file.Sync(); err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package wire

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"sync"
	"time"

	"github.com/jellevandenhooff/keytree/crypto"
)

var ErrNotFound = errors.New("not found")

func backoff(retries int) time.Duration {
	backoff, max := 1*time.Second, 10*time.Second
	for backoff < max && retries > 0 {
		backoff = backoff * 2
		retries--
	}
	if backoff > max {
		backoff = max
	}

	backoff -= time.Duration(float64(backoff) * 0.4 * rand.Float64())
	if backoff < 0 {
		return 0
	}
	return backoff
}

type client struct {
	host       string
	httpClient *http.Client

	mu      sync.Mutex
	retries int
}

func (c *client) success() {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.retries = 0
}

func (c *client) failure() {
	c.mu.Lock()
	defer c.mu.Unlock()
	// detect if we're already locked... if so, do nothing? (case of 8 simultaneous failures)

	time.Sleep(backoff(c.retries))
	c.retries += 1
}

func (c *client) await() {
	c.mu.Lock()
	defer c.mu.Unlock()
}

type KeyTreeClient struct {
	client *client
}

func decode(resp *http.Response, reply interface{}) error {
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return errors.New(resp.Status)
	}

	reader := io.LimitReader(resp.Body, 32*1024)
	dec := json.NewDecoder(reader)
	return dec.Decode(reply)
}

func (c *client) get(path string, reply interface{}) (err error) {
	c.await()

	defer func() {
		if err == nil {
			c.success()
		} else {
			c.failure()
		}
	}()

	resp, err := c.httpClient.Get(c.host + path)
	if err != nil {
		return err
	}

	return decode(resp, reply)
}

func (c *client) post(path string, request, reply interface{}) (err error) {
	c.await()

	defer func() {
		if err == nil {
			c.success()
		} else {
			c.failure()
		}
	}()

	// TODO: Use a pool of buffers?
	buffer := new(bytes.Buffer)
	if err := json.NewEncoder(buffer).Encode(request); err != nil {
		return err
	}

	resp, err := c.httpClient.Post(c.host+path, "text/json; charset=utf8", buffer)
	if err != nil {
		return err
	}

	return decode(resp, reply)
}

func newClient(host string) *client {
	return &client{
		host: host,
		httpClient: &http.Client{
			Timeout: 20 * time.Second,
		},
	}
}

func NewKeyTreeClient(host string) *KeyTreeClient {
	return &KeyTreeClient{client: newClient(host)}
}

func (c *KeyTreeClient) Submit(update *SignedEntry) error {
	var reply interface{}
	return c.client.post("/keytree/submit", update, &reply)
}

func (c *KeyTreeClient) TrieNode(h crypto.Hash) (*TrieNode, error) {
	var reply *TrieNode
	if err := c.client.get(fmt.Sprintf("/keytree/trienode?hash=%s", h), &reply); err != nil {
		return nil, err
	}
	if reply == nil {
		return nil, ErrNotFound
	}
	if err := reply.Check(); err != nil {
		return nil, err
	}
	return reply, nil
}

func (c *KeyTreeClient) Root() (*SignedRoot, error) {
	var reply SignedRoot
	if err := c.client.get("/keytree/root", &reply); err != nil {
		return nil, err
	}
	if err := reply.Check(); err != nil {
		return nil, err
	}
	return &reply, nil
}

func (c *KeyTreeClient) UpdateBatch(h crypto.Hash) (*UpdateBatch, error) {
	var reply *UpdateBatch
	if err := c.client.get(fmt.Sprintf("/keytree/updatebatch?hash=%s", h), &reply); err != nil {
		return nil, err
	}
	if reply == nil {
		return nil, ErrNotFound
	}
	if err := reply.Check(); err != nil {
		return nil, err
	}
	return reply, nil
}

func (c *KeyTreeClient) Lookup(h crypto.Hash) (*LookupReply, error) {
	var reply LookupReply
	if err := c.client.get(fmt.Sprintf("/keytree/lookup?hash=%s", h), &reply); err != nil {
		return nil, err
	}
	if err := reply.Check(); err != nil {
		return nil, err
	}
	return &reply, nil
}

func (c *KeyTreeClient) History(h crypto.Hash, since uint64) (*SignedEntry, error) {
	var reply *SignedEntry
	if err := c.client.get(fmt.Sprintf("/keytree/history?hash=%s&since=%d", h, since), &reply); err != nil {
		return nil, err
	}
	if reply == nil {
		return nil, ErrNotFound
	}
	if err := reply.Check(); err != nil {
		return nil, err
	}
	return reply, nil
}

type DKIMClient struct {
	client *client
}

func NewDKIMClient(host string) *DKIMClient {
	return &DKIMClient{client: newClient(host)}
}

func (c *DKIMClient) Prepare(req *DKIMStatement) (string, error) {
	var reply string
	if err := c.client.post("/dkim/prepare", req, &reply); err != nil {
		return "", err
	}
	return reply, nil
}

func (c *DKIMClient) Poll(req string) (*DKIMStatus, error) {
	var reply DKIMStatus
	if err := c.client.get(fmt.Sprintf("/dkim/poll?email=%s", req), &reply); err != nil {
		return nil, err
	}
	if err := reply.Check(); err != nil {
		return nil, err
	}
	return &reply, nil
}
<commit_msg>simplify a bit<commit_after>package wire

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"sync"
	"time"

	"github.com/jellevandenhooff/keytree/crypto"
)

var ErrNotFound = errors.New("not found")

func backoff(retries int) time.Duration {
	backoff := 1 * time.Second
	max := 60 * time.Second

	for i := 0; i < retries && backoff < max; i++ {
		backoff = backoff * 2
	}
	if backoff > max {
		backoff = max
	}

	backoff -= time.Duration(float64(backoff) * 0.4 * rand.Float64())
	if backoff < 0 {
		return 0
	}

	return backoff
}

type client struct {
	host       string
	httpClient *http.Client

	mu      sync.Mutex
	retries int
	waiting chan struct{}
}

func (c *client) success() {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.retries = 0
}

func (c *client) failure() {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.waiting != nil {
		return
	}

	c.waiting = make(chan struct{})
	time.AfterFunc(backoff(c.retries), func() {
		c.mu.Lock()
		defer c.mu.Unlock()

		close(c.waiting)
		c.waiting = nil
	})
	c.retries += 1
}

func (c *client) await() {
	c.mu.Lock()
	w := c.waiting
	c.mu.Unlock()

	if w != nil {
		<-w
	}
}

func (c *client) process(err *error) {
	if *err != nil {
		c.failure()
	} else {
		c.success()
	}
}

func decode(resp *http.Response, reply interface{}) error {
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return errors.New(resp.Status)
	}

	reader := io.LimitReader(resp.Body, 32*1024)
	dec := json.NewDecoder(reader)
	return dec.Decode(reply)
}

func (c *client) get(path string, reply interface{}) (err error) {
	c.await()
	defer c.process(&err)

	resp, err := c.httpClient.Get(c.host + path)
	if err != nil {
		return err
	}

	return decode(resp, reply)
}

func (c *client) post(path string, request, reply interface{}) (err error) {
	c.await()
	defer c.process(&err)

	buffer := new(bytes.Buffer)
	if err := json.NewEncoder(buffer).Encode(request); err != nil {
		return err
	}

	resp, err := c.httpClient.Post(c.host+path, "text/json; charset=utf8", buffer)
	if err != nil {
		return err
	}

	return decode(resp, reply)
}

func newClient(host string) *client {
	return &client{
		host: host,
		httpClient: &http.Client{
			Timeout: 20 * time.Second,
		},
	}
}

type KeyTreeClient struct {
	client *client
}

func NewKeyTreeClient(host string) *KeyTreeClient {
	return &KeyTreeClient{client: newClient(host)}
}

func (c *KeyTreeClient) Submit(update *SignedEntry) error {
	var reply interface{}
	return c.client.post("/keytree/submit", update, &reply)
}

func (c *KeyTreeClient) TrieNode(h crypto.Hash) (*TrieNode, error) {
	var reply *TrieNode
	if err := c.client.get(fmt.Sprintf("/keytree/trienode?hash=%s", h), &reply); err != nil {
		return nil, err
	}
	if reply == nil {
		return nil, ErrNotFound
	}
	if err := reply.Check(); err != nil {
		return nil, err
	}
	return reply, nil
}

func (c *KeyTreeClient) Root() (*SignedRoot, error) {
	var reply SignedRoot
	if err := c.client.get("/keytree/root", &reply); err != nil {
		return nil, err
	}
	if err := reply.Check(); err != nil {
		return nil, err
	}
	return &reply, nil
}

func (c *KeyTreeClient) UpdateBatch(h crypto.Hash) (*UpdateBatch, error) {
	var reply *UpdateBatch
	if err := c.client.get(fmt.Sprintf("/keytree/updatebatch?hash=%s", h), &reply); err != nil {
		return nil, err
	}
	if reply == nil {
		return nil, ErrNotFound
	}
	if err := reply.Check(); err != nil {
		return nil, err
	}
	return reply, nil
}

func (c *KeyTreeClient) Lookup(h crypto.Hash) (*LookupReply, error) {
	var reply LookupReply
	if err := c.client.get(fmt.Sprintf("/keytree/lookup?hash=%s", h), &reply); err != nil {
		return nil, err
	}
	if err := reply.Check(); err != nil {
		return nil, err
	}
	return &reply, nil
}

func (c *KeyTreeClient) History(h crypto.Hash, since uint64) (*SignedEntry, error) {
	var reply *SignedEntry
	if err := c.client.get(fmt.Sprintf("/keytree/history?hash=%s&since=%d", h, since), &reply); err != nil {
		return nil, err
	}
	if reply == nil {
		return nil, ErrNotFound
	}
	if err := reply.Check(); err != nil {
		return nil, err
	}
	return reply, nil
}

type DKIMClient struct {
	client *client
}

func NewDKIMClient(host string) *DKIMClient {
	return &DKIMClient{client: newClient(host)}
}

func (c *DKIMClient) Prepare(req *DKIMStatement) (string, error) {
	var reply string
	if err := c.client.post("/dkim/prepare", req, &reply); err != nil {
		return "", err
	}
	return reply, nil
}

func (c *DKIMClient) Poll(req string) (*DKIMStatus, error) {
	var reply DKIMStatus
	if err := c.client.get(fmt.Sprintf("/dkim/poll?email=%s", req), &reply); err != nil {
		return nil, err
	}
	if err := reply.Check(); err != nil {
		return nil, err
	}
	return &reply, nil
}
<|endoftext|>


<commit_before>package dnsd

import (
	"context"
	"math/rand"
	"net"
	"strings"
	"time"

	"github.com/HouzuoGuo/laitos/lalog"
	"github.com/HouzuoGuo/laitos/toolbox"
	"golang.org/x/net/dns/dnsmessage"

	"github.com/HouzuoGuo/laitos/misc"
)

// GetTCPStatsCollector returns stats collector for the TCP server of this daemon.
func (daemon *Daemon) GetTCPStatsCollector() *misc.Stats {
	return misc.DNSDStatsTCP
}

// HandleTCPConnection reads a DNS query from a TCP client and responds to it with the DNS query result.
func (daemon *Daemon) HandleTCPConnection(logger lalog.Logger, ip string, conn *net.TCPConn) {
	misc.TweakTCPConnection(conn, ClientTimeoutSec*time.Second)
	// Read query length
	queryLen := make([]byte, 2)
	_, err := conn.Read(queryLen)
	if err != nil {
		logger.Warning("HandleTCPConnection", ip, err, "failed to read query length from client")
		return
	}
	queryLenInteger := int(queryLen[0])*256 + int(queryLen[1])
	// Read query packet
	if queryLenInteger > MaxPacketSize || queryLenInteger < MinNameQuerySize {
		logger.Info("HandleTCPConnection", ip, nil, "invalid query length from client")
		return
	}
	queryBody := make([]byte, queryLenInteger)
	_, err = conn.Read(queryBody)
	if err != nil {
		logger.Warning("HandleTCPConnection", ip, err, "failed to read query from client")
		return
	}
	// Parse the first (and only) query question.
	parser := new(dnsmessage.Parser)
	header, err := parser.Start(queryBody)
	if err != nil {
		logger.Warning("HandleTCPConnection", ip, err, "failed to parse query header")
		return
	}
	question, err := parser.Question()
	if err != nil {
		logger.Warning("HandleTCPConnection", ip, err, "failed to parse query question")
		return
	}
	var respBody []byte
	if question.Type == dnsmessage.TypeTXT {
		// The TXT query may be carrying an app command.
		respBody = daemon.handleTCPTextQuery(ip, queryLen, queryBody, header, question)
	} else {
		// Handle all other query types.
		respBody = daemon.handleTCPNameOrOtherQuery(ip, queryLen, queryBody, header, question)
	}
	// Return early (and close the client connection) in case there is no
	// appropriate response.
	if len(respBody) < 3 {
		return
	}
	// Match the response transaction ID with the request.
	respBody[0] = queryBody[0]
	respBody[1] = queryBody[1]
	// Reset connection IO timeout.
	misc.TweakTCPConnection(conn, ClientTimeoutSec*time.Second)
	respLen := []byte{byte(len(respBody) / 256), byte(len(respBody) % 256)}
	if _, err := conn.Write(respLen); err != nil {
		logger.Warning("HandleTCPConnection", ip, err, "failed to answer length to the client")
		return
	} else if _, err := conn.Write(respBody); err != nil {
		logger.Warning("HandleTCPConnection", ip, err, "failed to answer to the client")
		return
	}
}

func (daemon *Daemon) handleTCPTextQuery(clientIP string, queryLen, queryBody []byte, header dnsmessage.Header, question dnsmessage.Question) (respBody []byte) {
	name := question.Name.String()
	if daemon.processQueryTestCaseFunc != nil {
		daemon.processQueryTestCaseFunc(name)
	}
	// Remove the domain name from the labels.
	labelsWithoutDomain := strings.Split(name, ".")
	if len(labelsWithoutDomain) > 2 {
		labelsWithoutDomain = labelsWithoutDomain[:len(labelsWithoutDomain)-2]
	}
	if dtmfDecoded := DecodeDTMFCommandInput(labelsWithoutDomain); len(dtmfDecoded) > 1 {
		cmdResult := daemon.latestCommands.Execute(context.Background(), daemon.Processor, clientIP, dtmfDecoded)
		if cmdResult.Error == toolbox.ErrPINAndShortcutNotFound {
			// Because the prefix may appear in an ordinary text record query
			// that is not a toolbox command, when there is a PIN mismatch,
			// forward to recursive resolver as if the query is indeed not a
			// toolbox command.
			daemon.logger.Info("handleTCPTextQuery", clientIP, nil, "the queried name has the command prefix but failed PIN check, forwarding to recursive resolver.")
			goto forwardToRecursiveResolver
		} else {
			var err error
			daemon.logger.Info("handleTCPTextQuery", clientIP, nil, "processed a toolbox command")
			// Try to fit the response into a single TXT entry.
			// Keep in mind that by convention DNS uses 512 bytes as the overall
			// message size limit - including both question and response.
			// Leave some buffer room for the DNS headers.
			respBody, err = BuildTextResponse(name, header, question, misc.SplitIntoSlice(cmdResult.CombinedOutput, 200, 200))
			if err != nil {
				daemon.logger.Warning("handleTCPTextQuery", clientIP, err, "failed to build response packet")
				return nil
			}
			return
		}
	} else {
		daemon.logger.Info("handleTCPTextQuery", clientIP, nil, "handle TXT query %q", name)
	}
forwardToRecursiveResolver:
	// Because the password PIN in the query may contain a typo, make sure this
	// function does not incidentally log the query name.
	return daemon.handleTCPRecursiveQuery(clientIP, queryLen, queryBody)
}

func (daemon *Daemon) handleTCPNameOrOtherQuery(clientIP string, queryLen, queryBody []byte, header dnsmessage.Header, question dnsmessage.Question) (respBody []byte) {
	name := question.Name.String()
	if name == "" {
		daemon.logger.Info("handleTCPNameOrOtherQuery", clientIP, nil, "handle non-name query")
	} else {
		if daemon.processQueryTestCaseFunc != nil {
			daemon.processQueryTestCaseFunc(name)
		}
		daemon.logger.Info("handleTCPNameOrOtherQuery", clientIP, nil, "handle name query %q", name)
	}
	if daemon.IsInBlacklist(name) {
		daemon.logger.Info("handleTCPNameOrOtherQuery", clientIP, nil, "handle black-listed name query %q", name)
		var err error
		respBody, err = BuildBlackHoleAddrResponse(header, question)
		if err != nil {
			daemon.logger.Warning("handleTCPNameOrOtherQuery", clientIP, err, "failed to build response packet")
			return nil
		}
	} else {
		respBody = daemon.handleTCPRecursiveQuery(clientIP, queryLen, queryBody)
	}
	return
}

/*
handleTCPRecursiveQuery forward the input query to a randomly chosen recursive resolver and retrieves the response.
Be aware that toolbox command processor may invoke this function with an incorrect PIN entry similar to the real PIN,
therefore this function must not log the input packet content in any way.
*/
func (daemon *Daemon) handleTCPRecursiveQuery(clientIP string, queryLen, queryBody []byte) (respBody []byte) {
	respBody = make([]byte, 0)
	if !daemon.checkAllowClientIP(clientIP) {
		daemon.logger.Warning("handleTCPRecursiveQuery", clientIP, nil, "client IP is not allowed to query")
		return
	}
	randForwarder := daemon.Forwarders[rand.Intn(len(daemon.Forwarders))]
	// Forward the query to a randomly chosen recursive resolver
	myForwarder, err := net.DialTimeout("tcp", randForwarder, ForwarderTimeoutSec*time.Second)
	if err != nil {
		daemon.logger.Warning("handleTCPRecursiveQuery", clientIP, err, "failed to connect to forwarder")
		return
	}
	defer func() {
		daemon.logger.MaybeMinorError(myForwarder.Close())
	}()
	// Send original query to the resolver without modification
	daemon.logger.MaybeMinorError(myForwarder.SetDeadline(time.Now().Add(ForwarderTimeoutSec * time.Second)))
	if _, err = myForwarder.Write(queryLen); err != nil {
		daemon.logger.Warning("handleTCPRecursiveQuery", clientIP, err, "failed to write length to forwarder")
		return
	} else if _, err = myForwarder.Write(queryBody); err != nil {
		daemon.logger.Warning("handleTCPRecursiveQuery", clientIP, err, "failed to write query to forwarder")
		return
	}
	// Read resolver's response
	respLen := make([]byte, 2)
	if _, err = myForwarder.Read(respLen); err != nil {
		daemon.logger.Warning("handleTCPRecursiveQuery", clientIP, err, "failed to read length from forwarder")
		return
	}
	respLenInt := int(respLen[0])*256 + int(respLen[1])
	if respLenInt > MaxPacketSize || respLenInt < 1 {
		daemon.logger.Warning("handleTCPRecursiveQuery", clientIP, nil, "bad response length from forwarder")
		return
	}
	respBody = make([]byte, respLenInt)
	if _, err = myForwarder.Read(respBody); err != nil {
		daemon.logger.Warning("handleTCPRecursiveQuery", clientIP, err, "failed to read response from forwarder")
		return
	}
	return
}
<commit_msg>Log a couple of local variables to debug test case failure on circleci.<commit_after>package dnsd

import (
	"context"
	"math/rand"
	"net"
	"strings"
	"time"

	"github.com/HouzuoGuo/laitos/lalog"
	"github.com/HouzuoGuo/laitos/toolbox"
	"golang.org/x/net/dns/dnsmessage"

	"github.com/HouzuoGuo/laitos/misc"
)

// GetTCPStatsCollector returns stats collector for the TCP server of this daemon.
func (daemon *Daemon) GetTCPStatsCollector() *misc.Stats {
	return misc.DNSDStatsTCP
}

// HandleTCPConnection reads a DNS query from a TCP client and responds to it with the DNS query result.
func (daemon *Daemon) HandleTCPConnection(logger lalog.Logger, ip string, conn *net.TCPConn) {
	misc.TweakTCPConnection(conn, ClientTimeoutSec*time.Second)
	// Read query length
	queryLen := make([]byte, 2)
	_, err := conn.Read(queryLen)
	if err != nil {
		logger.Warning("HandleTCPConnection", ip, err, "failed to read query length from client")
		return
	}
	queryLenInteger := int(queryLen[0])*256 + int(queryLen[1])
	// Read query packet
	if queryLenInteger > MaxPacketSize || queryLenInteger < MinNameQuerySize {
		logger.Info("HandleTCPConnection", ip, nil, "invalid query length from client")
		return
	}
	queryBody := make([]byte, queryLenInteger)
	_, err = conn.Read(queryBody)
	if err != nil {
		logger.Warning("HandleTCPConnection", ip, err, "failed to read query from client")
		return
	}
	// Parse the first (and only) query question.
	parser := new(dnsmessage.Parser)
	header, err := parser.Start(queryBody)
	if err != nil {
		logger.Warning("HandleTCPConnection", ip, err, "failed to parse query header")
		return
	}
	question, err := parser.Question()
	if err != nil {
		logger.Warning("HandleTCPConnection", ip, err, "failed to parse query question")
		return
	}
	var respBody []byte
	if question.Type == dnsmessage.TypeTXT {
		// The TXT query may be carrying an app command.
		respBody = daemon.handleTCPTextQuery(ip, queryLen, queryBody, header, question)
	} else {
		// Handle all other query types.
		respBody = daemon.handleTCPNameOrOtherQuery(ip, queryLen, queryBody, header, question)
	}
	// Return early (and close the client connection) in case there is no
	// appropriate response.
	if len(respBody) < 3 {
		return
	}
	// Match the response transaction ID with the request.
	respBody[0] = queryBody[0]
	respBody[1] = queryBody[1]
	// Reset connection IO timeout.
	misc.TweakTCPConnection(conn, ClientTimeoutSec*time.Second)
	respLen := []byte{byte(len(respBody) / 256), byte(len(respBody) % 256)}
	if _, err := conn.Write(respLen); err != nil {
		logger.Warning("HandleTCPConnection", ip, err, "failed to answer length to the client")
		return
	} else if _, err := conn.Write(respBody); err != nil {
		logger.Warning("HandleTCPConnection", ip, err, "failed to answer to the client")
		return
	}
}

func (daemon *Daemon) handleTCPTextQuery(clientIP string, queryLen, queryBody []byte, header dnsmessage.Header, question dnsmessage.Question) (respBody []byte) {
	name := question.Name.String()
	if daemon.processQueryTestCaseFunc != nil {
		daemon.processQueryTestCaseFunc(name)
	}
	lalog.DefaultLogger.Warning("@@@@@@@@@@@@@@@@@@", "", nil, "input name: %q", name)
	// Remove the domain name from the labels.
	labelsWithoutDomain := strings.Split(name, ".")
	if len(labelsWithoutDomain) > 2 {
		labelsWithoutDomain = labelsWithoutDomain[:len(labelsWithoutDomain)-2]
	}
	lalog.DefaultLogger.Warning("@@@@@@@@@@@@@@@@@@", "", nil, "labelsWithoutDomain: %+v", labelsWithoutDomain)
	if dtmfDecoded := DecodeDTMFCommandInput(labelsWithoutDomain); len(dtmfDecoded) > 1 {
		lalog.DefaultLogger.Warning("@@@@@@@@@@@@@@@@@@", "", nil, "dtmfDecoded: %q", dtmfDecoded)
		cmdResult := daemon.latestCommands.Execute(context.Background(), daemon.Processor, clientIP, dtmfDecoded)
		if cmdResult.Error == toolbox.ErrPINAndShortcutNotFound {
			// Because the prefix may appear in an ordinary text record query
			// that is not a toolbox command, when there is a PIN mismatch,
			// forward to recursive resolver as if the query is indeed not a
			// toolbox command.
			daemon.logger.Info("handleTCPTextQuery", clientIP, nil, "the queried name has the command prefix but failed PIN check, forwarding to recursive resolver.")
			goto forwardToRecursiveResolver
		} else {
			var err error
			daemon.logger.Info("handleTCPTextQuery", clientIP, nil, "processed a toolbox command")
			// Try to fit the response into a single TXT entry.
			// Keep in mind that by convention DNS uses 512 bytes as the overall
			// message size limit - including both question and response.
			// Leave some buffer room for the DNS headers.
			respBody, err = BuildTextResponse(name, header, question, misc.SplitIntoSlice(cmdResult.CombinedOutput, 200, 200))
			if err != nil {
				daemon.logger.Warning("handleTCPTextQuery", clientIP, err, "failed to build response packet")
				return nil
			}
			return
		}
	} else {
		daemon.logger.Info("handleTCPTextQuery", clientIP, nil, "handle TXT query %q", name)
	}
forwardToRecursiveResolver:
	// Because the password PIN in the query may contain a typo, make sure this
	// function does not incidentally log the query name.
	return daemon.handleTCPRecursiveQuery(clientIP, queryLen, queryBody)
}

func (daemon *Daemon) handleTCPNameOrOtherQuery(clientIP string, queryLen, queryBody []byte, header dnsmessage.Header, question dnsmessage.Question) (respBody []byte) {
	name := question.Name.String()
	if name == "" {
		daemon.logger.Info("handleTCPNameOrOtherQuery", clientIP, nil, "handle non-name query")
	} else {
		if daemon.processQueryTestCaseFunc != nil {
			daemon.processQueryTestCaseFunc(name)
		}
		daemon.logger.Info("handleTCPNameOrOtherQuery", clientIP, nil, "handle name query %q", name)
	}
	if daemon.IsInBlacklist(name) {
		daemon.logger.Info("handleTCPNameOrOtherQuery", clientIP, nil, "handle black-listed name query %q", name)
		var err error
		respBody, err = BuildBlackHoleAddrResponse(header, question)
		if err != nil {
			daemon.logger.Warning("handleTCPNameOrOtherQuery", clientIP, err, "failed to build response packet")
			return nil
		}
	} else {
		respBody = daemon.handleTCPRecursiveQuery(clientIP, queryLen, queryBody)
	}
	return
}

/*
handleTCPRecursiveQuery forward the input query to a randomly chosen recursive resolver and retrieves the response.
Be aware that toolbox command processor may invoke this function with an incorrect PIN entry similar to the real PIN,
therefore this function must not log the input packet content in any way.
*/
func (daemon *Daemon) handleTCPRecursiveQuery(clientIP string, queryLen, queryBody []byte) (respBody []byte) {
	respBody = make([]byte, 0)
	if !daemon.checkAllowClientIP(clientIP) {
		daemon.logger.Warning("handleTCPRecursiveQuery", clientIP, nil, "client IP is not allowed to query")
		return
	}
	randForwarder := daemon.Forwarders[rand.Intn(len(daemon.Forwarders))]
	// Forward the query to a randomly chosen recursive resolver
	myForwarder, err := net.DialTimeout("tcp", randForwarder, ForwarderTimeoutSec*time.Second)
	if err != nil {
		daemon.logger.Warning("handleTCPRecursiveQuery", clientIP, err, "failed to connect to forwarder")
		return
	}
	defer func() {
		daemon.logger.MaybeMinorError(myForwarder.Close())
	}()
	// Send original query to the resolver without modification
	daemon.logger.MaybeMinorError(myForwarder.SetDeadline(time.Now().Add(ForwarderTimeoutSec * time.Second)))
	if _, err = myForwarder.Write(queryLen); err != nil {
		daemon.logger.Warning("handleTCPRecursiveQuery", clientIP, err, "failed to write length to forwarder")
		return
	} else if _, err = myForwarder.Write(queryBody); err != nil {
		daemon.logger.Warning("handleTCPRecursiveQuery", clientIP, err, "failed to write query to forwarder")
		return
	}
	// Read resolver's response
	respLen := make([]byte, 2)
	if _, err = myForwarder.Read(respLen); err != nil {
		daemon.logger.Warning("handleTCPRecursiveQuery", clientIP, err, "failed to read length from forwarder")
		return
	}
	respLenInt := int(respLen[0])*256 + int(respLen[1])
	if respLenInt > MaxPacketSize || respLenInt < 1 {
		daemon.logger.Warning("handleTCPRecursiveQuery", clientIP, nil, "bad response length from forwarder")
		return
	}
	respBody = make([]byte, respLenInt)
	if _, err = myForwarder.Read(respBody); err != nil {
		daemon.logger.Warning("handleTCPRecursiveQuery", clientIP, err, "failed to read response from forwarder")
		return
	}
	return
}
<|endoftext|>


<commit_before>/*
   Copyright The containerd 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 cgroups

import (
	"strings"
	"testing"

	v1 "github.com/containerd/cgroups/stats/v1"
)

const data = `   7       0 loop0 0 0 0 0 0 0 0 0 0 0 0
   7       1 loop1 0 0 0 0 0 0 0 0 0 0 0
   7       2 loop2 0 0 0 0 0 0 0 0 0 0 0
   7       3 loop3 0 0 0 0 0 0 0 0 0 0 0
   7       4 loop4 0 0 0 0 0 0 0 0 0 0 0
   7       5 loop5 0 0 0 0 0 0 0 0 0 0 0
   7       6 loop6 0 0 0 0 0 0 0 0 0 0 0
   7       7 loop7 0 0 0 0 0 0 0 0 0 0 0
   8       0 sda 1892042 187697 63489222 1246284 1389086 2887005 134903104 11390608 1 1068060 12692228
   8       1 sda1 1762875 37086 61241570 1200512 1270037 2444415 131214808 11152764 1 882624 12409308
   8       2 sda2 2 0 4 0 0 0 0 0 0 0 0
   8       5 sda5 129102 150611 2244440 45716 18447 442590 3688296 67268 0 62584 112984`

func TestGetDevices(t *testing.T) {
	r := strings.NewReader(data)
	devices, err := getDevices(r)
	if err != nil {
		t.Fatal(err)
	}
	name, ok := devices[deviceKey{8, 0}]
	if !ok {
		t.Fatal("no device found for 8,0")
	}
	const expected = "/dev/sda"
	if name != expected {
		t.Fatalf("expected device name %q but received %q", expected, name)
	}
}

func TestNewBlkio(t *testing.T) {
	const root = "/test/folder"
	const expected = "/test/folder/blkio"
	const expectedProc = "/proc"

	ctrl := NewBlkio(root)
	if ctrl.root != expected {
		t.Fatalf("expected cgroups root %q but received %q", expected, ctrl.root)
	}
	if ctrl.procRoot != expectedProc {
		t.Fatalf("expected proc FS root %q but received %q", expectedProc, ctrl.procRoot)
	}
}

func TestBlkioStat(t *testing.T) {
	ctrl := NewBlkio("/sys/fs/cgroup")

	var metrics v1.Metrics
	err := ctrl.Stat("", &metrics)
	if err != nil {
		t.Fatalf("failed to call Stat: %v", err)
	}

	if len(metrics.Blkio.IoServicedRecursive) == 0 {
		t.Fatalf("IoServicedRecursive must not be empty")
	}
	if len(metrics.Blkio.IoServiceBytesRecursive) == 0 {
		t.Fatalf("IoServiceBytesRecursive must not be empty")
	}
}

func TestNewBlkio_Proc(t *testing.T) {
	const root = "/test/folder"
	const expected = "/test/folder/blkio"
	const expectedProc = "/test/proc"

	ctrl := NewBlkio(root, ProcRoot(expectedProc))
	if ctrl.root != expected {
		t.Fatalf("expected cgroups root %q but received %q", expected, ctrl.root)
	}
	if ctrl.procRoot != expectedProc {
		t.Fatalf("expected proc FS root %q but received %q", expectedProc, ctrl.procRoot)
	}
}
<commit_msg>Skip TestBlkioStat() if there is no /sys/fs/cgroup/blkio<commit_after>/*
   Copyright The containerd 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 cgroups

import (
	"os"
	"strings"
	"testing"

	v1 "github.com/containerd/cgroups/stats/v1"
)

const data = `   7       0 loop0 0 0 0 0 0 0 0 0 0 0 0
   7       1 loop1 0 0 0 0 0 0 0 0 0 0 0
   7       2 loop2 0 0 0 0 0 0 0 0 0 0 0
   7       3 loop3 0 0 0 0 0 0 0 0 0 0 0
   7       4 loop4 0 0 0 0 0 0 0 0 0 0 0
   7       5 loop5 0 0 0 0 0 0 0 0 0 0 0
   7       6 loop6 0 0 0 0 0 0 0 0 0 0 0
   7       7 loop7 0 0 0 0 0 0 0 0 0 0 0
   8       0 sda 1892042 187697 63489222 1246284 1389086 2887005 134903104 11390608 1 1068060 12692228
   8       1 sda1 1762875 37086 61241570 1200512 1270037 2444415 131214808 11152764 1 882624 12409308
   8       2 sda2 2 0 4 0 0 0 0 0 0 0 0
   8       5 sda5 129102 150611 2244440 45716 18447 442590 3688296 67268 0 62584 112984`

func TestGetDevices(t *testing.T) {
	r := strings.NewReader(data)
	devices, err := getDevices(r)
	if err != nil {
		t.Fatal(err)
	}
	name, ok := devices[deviceKey{8, 0}]
	if !ok {
		t.Fatal("no device found for 8,0")
	}
	const expected = "/dev/sda"
	if name != expected {
		t.Fatalf("expected device name %q but received %q", expected, name)
	}
}

func TestNewBlkio(t *testing.T) {
	const root = "/test/folder"
	const expected = "/test/folder/blkio"
	const expectedProc = "/proc"

	ctrl := NewBlkio(root)
	if ctrl.root != expected {
		t.Fatalf("expected cgroups root %q but received %q", expected, ctrl.root)
	}
	if ctrl.procRoot != expectedProc {
		t.Fatalf("expected proc FS root %q but received %q", expectedProc, ctrl.procRoot)
	}
}

func TestBlkioStat(t *testing.T) {
	_, err := os.Stat("/sys/fs/cgroup/blkio")
	if os.IsNotExist(err) {
		t.Skip("failed to find /sys/fs/cgroup/blkio")
	}

	ctrl := NewBlkio("/sys/fs/cgroup")

	var metrics v1.Metrics
	err = ctrl.Stat("", &metrics)
	if err != nil {
		t.Fatalf("failed to call Stat: %v", err)
	}

	if len(metrics.Blkio.IoServicedRecursive) == 0 {
		t.Fatalf("IoServicedRecursive must not be empty")
	}
	if len(metrics.Blkio.IoServiceBytesRecursive) == 0 {
		t.Fatalf("IoServiceBytesRecursive must not be empty")
	}
}

func TestNewBlkio_Proc(t *testing.T) {
	const root = "/test/folder"
	const expected = "/test/folder/blkio"
	const expectedProc = "/test/proc"

	ctrl := NewBlkio(root, ProcRoot(expectedProc))
	if ctrl.root != expected {
		t.Fatalf("expected cgroups root %q but received %q", expected, ctrl.root)
	}
	if ctrl.procRoot != expectedProc {
		t.Fatalf("expected proc FS root %q but received %q", expectedProc, ctrl.procRoot)
	}
}
<|endoftext|>


<commit_before>package app

import (
	"github.com/revel/revel"
	"goassimp/app/mgnredis"
	"time"
)

func init() {
	// Filters is the default set of global filters.
	revel.Filters = []revel.Filter{
		revel.PanicFilter,             // Recover from panics and display an error page instead.
		revel.RouterFilter,            // Use the routing table to select the right Action
		revel.FilterConfiguringFilter, // A hook for adding or removing per-Action filters.
		revel.ParamsFilter,            // Parse parameters into Controller.Params.
		revel.SessionFilter,           // Restore and write the session cookie.
		revel.FlashFilter,             // Restore and write the flash cookie.
		revel.ValidationFilter,        // Restore kept validation errors and save new ones from cookie.
		revel.I18nFilter,              // Resolve the requested language
		HeaderFilter,                  // Add some security based headers
		revel.InterceptorFilter,       // Run interceptors around the action.
		revel.CompressFilter,          // Compress the result.
		revel.ActionInvoker,           // Invoke the action.
	}

	// register startup functions with OnAppStart
	// ( order dependent )
	revel.OnAppStart(func() {
		host := revel.Config.StringDefault("redis.host", ":6379")
		capacity := revel.Config.IntDefault("redis.capacity_pool", 20)
		maxcapacity := revel.Config.IntDefault("redis.max_capacity_pool", 200)
		idleTimeout := revel.Config.StringDefault("redis.idleTimeout", "10m")
		duration, err := time.ParseDuration(idleTimeout)
		if err != nil {
			panic(err)
		}
		mgnredis.InitRedis(host, capacity, maxcapacity, duration)
	})
}

// TODO turn this into revel.HeaderFilter
// should probably also have a filter for CSRF
// not sure if it can go in the same filter or not
var HeaderFilter = func(c *revel.Controller, fc []revel.Filter) {
	// Add some common security headers
	c.Response.Out.Header().Add("X-Frame-Options", "SAMEORIGIN")
	c.Response.Out.Header().Add("X-XSS-Protection", "1; mode=block")
	c.Response.Out.Header().Add("X-Content-Type-Options", "nosniff")

	fc[0](c, fc[1:]) // Execute the next filter stage.
}
<commit_msg>redis from comf<commit_after>package app

import (
	"github.com/revel/revel"
	"goassimp/app/mgnredis"
	"time"
)

func init() {
	// Filters is the default set of global filters.
	revel.Filters = []revel.Filter{
		revel.PanicFilter,             // Recover from panics and display an error page instead.
		revel.RouterFilter,            // Use the routing table to select the right Action
		revel.FilterConfiguringFilter, // A hook for adding or removing per-Action filters.
		revel.ParamsFilter,            // Parse parameters into Controller.Params.
		revel.SessionFilter,           // Restore and write the session cookie.
		revel.FlashFilter,             // Restore and write the flash cookie.
		revel.ValidationFilter,        // Restore kept validation errors and save new ones from cookie.
		revel.I18nFilter,              // Resolve the requested language
		HeaderFilter,                  // Add some security based headers
		revel.InterceptorFilter,       // Run interceptors around the action.
		revel.CompressFilter,          // Compress the result.
		revel.ActionInvoker,           // Invoke the action.
	}

	// register startup functions with OnAppStart
	// ( order dependent )
	revel.OnAppStart(func() {
		host := revel.Config.StringDefault("redis.host", ":6379")
		capacity := revel.Config.IntDefault("redis.capacity_pool", 20)
		max_capacity := revel.Config.IntDefault("redis.max_capacity_pool", 200)
		dstr := revel.Config.StringDefault("redis.idleTimeout", "1m")
		duration, err := time.ParseDuration(dstr)
		if err != nil {
			panic(err)
		}
		mgnredis.InitRedis(host, capacity, max_capacity, duration)
	})
}

// TODO turn this into revel.HeaderFilter
// should probably also have a filter for CSRF
// not sure if it can go in the same filter or not
var HeaderFilter = func(c *revel.Controller, fc []revel.Filter) {
	// Add some common security headers
	c.Response.Out.Header().Add("X-Frame-Options", "SAMEORIGIN")
	c.Response.Out.Header().Add("X-XSS-Protection", "1; mode=block")
	c.Response.Out.Header().Add("X-Content-Type-Options", "nosniff")

	fc[0](c, fc[1:]) // Execute the next filter stage.
}
<|endoftext|>


<commit_before>package http

import (
	"bytes"
	"io"
	"net"
	"net/http"
	"time"

	"v2ray.com/core/common/alloc"
	"v2ray.com/core/common/loader"
	"v2ray.com/core/transport/internet"
)

const (
	CRLF   = "\r\n"
	ENDING = CRLF + CRLF
)

type Reader interface {
	Read(io.Reader) (*alloc.Buffer, error)
}

type Writer interface {
	Write(io.Writer) error
}

type NoOpReader struct{}

func (this *NoOpReader) Read(io.Reader) (*alloc.Buffer, error) {
	return nil, nil
}

type NoOpWriter struct{}

func (this *NoOpWriter) Write(io.Writer) error {
	return nil
}

type HeaderReader struct {
}

func (*HeaderReader) Read(reader io.Reader) (*alloc.Buffer, error) {
	buffer := alloc.NewSmallBuffer()
	for {
		_, err := buffer.FillFrom(reader)
		if err != nil {
			return nil, err
		}
		if n := bytes.Index(buffer.Value, []byte(ENDING)); n != -1 {
			buffer.SliceFrom(n + len(ENDING))
			break
		}
		if buffer.Len() >= len(ENDING) {
			copy(buffer.Value, buffer.Value[buffer.Len()-len(ENDING):])
			buffer.Slice(0, len(ENDING))
		}
	}
	if buffer.IsEmpty() {
		buffer.Release()
		return nil, nil
	}
	return buffer, nil
}

type HeaderWriter struct {
	header *alloc.Buffer
}

func NewHeaderWriter(header *alloc.Buffer) *HeaderWriter {
	return &HeaderWriter{
		header: header,
	}
}

func (this *HeaderWriter) Write(writer io.Writer) error {
	if this.header == nil {
		return nil
	}
	_, err := writer.Write(this.header.Value)
	this.header.Release()
	this.header = nil
	return err
}

type HttpConn struct {
	net.Conn

	readBuffer    *alloc.Buffer
	oneTimeReader Reader
	oneTimeWriter Writer
}

func NewHttpConn(conn net.Conn, reader Reader, writer Writer) *HttpConn {
	return &HttpConn{
		Conn:          conn,
		oneTimeReader: reader,
		oneTimeWriter: writer,
	}
}

func (this *HttpConn) Read(b []byte) (int, error) {
	if this.oneTimeReader != nil {
		buffer, err := this.oneTimeReader.Read(this.Conn)
		if err != nil {
			return 0, err
		}
		this.readBuffer = buffer
		this.oneTimeReader = nil
	}

	if this.readBuffer.Len() > 0 {
		nBytes, err := this.readBuffer.Read(b)
		if nBytes == this.readBuffer.Len() {
			this.readBuffer.Release()
			this.readBuffer = nil
		}
		return nBytes, err
	}

	return this.Conn.Read(b)
}

func (this *HttpConn) Write(b []byte) (int, error) {
	if this.oneTimeWriter != nil {
		err := this.oneTimeWriter.Write(this.Conn)
		this.oneTimeWriter = nil
		if err != nil {
			return 0, err
		}
	}

	return this.Conn.Write(b)
}

type HttpAuthenticator struct {
	config *Config
}

func (this HttpAuthenticator) GetClientWriter() *HeaderWriter {
	header := alloc.NewSmallBuffer().Clear()
	config := this.config.Request
	header.AppendString(config.Method.GetValue()).AppendString(" ").AppendString(config.PickUri()).AppendString(" ").AppendString(config.GetFullVersion()).AppendString(CRLF)

	headers := config.PickHeaders()
	for _, h := range headers {
		header.AppendString(h).AppendString(CRLF)
	}
	header.AppendString(CRLF)
	return &HeaderWriter{
		header: header,
	}
}

func (this HttpAuthenticator) GetServerWriter() *HeaderWriter {
	header := alloc.NewSmallBuffer().Clear()
	config := this.config.Response
	header.AppendString(config.GetFullVersion()).AppendString(" ").AppendString(config.Status.GetCode()).AppendString(" ").AppendString(config.Status.GetReason()).AppendString(CRLF)

	headers := config.PickHeaders()
	for _, h := range headers {
		header.AppendString(h).AppendString(CRLF)
	}
	if !config.HasHeader("Date") {
		header.AppendString("Date: ").AppendString(time.Now().Format(http.TimeFormat)).AppendString(CRLF)
	}
	header.AppendString(CRLF)
	return &HeaderWriter{
		header: header,
	}
}

func (this HttpAuthenticator) Client(conn net.Conn) net.Conn {
	if this.config.Request == nil && this.config.Response == nil {
		return conn
	}
	var reader Reader = new(NoOpReader)
	if this.config.Request != nil {
		reader = new(HeaderReader)
	}

	var writer Writer = new(NoOpWriter)
	if this.config.Response != nil {
		writer = this.GetClientWriter()
	}
	return NewHttpConn(conn, reader, writer)
}

func (this HttpAuthenticator) Server(conn net.Conn) net.Conn {
	if this.config.Request == nil && this.config.Response == nil {
		return conn
	}
	return NewHttpConn(conn, new(HeaderReader), this.GetServerWriter())
}

type HttpAuthenticatorFactory struct{}

func (HttpAuthenticatorFactory) Create(config interface{}) internet.ConnectionAuthenticator {
	return HttpAuthenticator{
		config: config.(*Config),
	}
}

func init() {
	internet.RegisterConnectionAuthenticator(loader.GetType(new(Config)), HttpAuthenticatorFactory{})
}
<commit_msg>fix http read<commit_after>package http

import (
	"bytes"
	"io"
	"net"
	"net/http"
	"time"

	"v2ray.com/core/common/alloc"
	"v2ray.com/core/common/loader"
	"v2ray.com/core/transport/internet"
)

const (
	CRLF   = "\r\n"
	ENDING = CRLF + CRLF
)

type Reader interface {
	Read(io.Reader) (*alloc.Buffer, error)
}

type Writer interface {
	Write(io.Writer) error
}

type NoOpReader struct{}

func (this *NoOpReader) Read(io.Reader) (*alloc.Buffer, error) {
	return nil, nil
}

type NoOpWriter struct{}

func (this *NoOpWriter) Write(io.Writer) error {
	return nil
}

type HeaderReader struct {
}

func (*HeaderReader) Read(reader io.Reader) (*alloc.Buffer, error) {
	buffer := alloc.NewSmallBuffer().Clear()
	for {
		_, err := buffer.FillFrom(reader)
		if err != nil {
			return nil, err
		}
		if n := bytes.Index(buffer.Value, []byte(ENDING)); n != -1 {
			buffer.SliceFrom(n + len(ENDING))
			break
		}
		if buffer.Len() >= len(ENDING) {
			copy(buffer.Value, buffer.Value[buffer.Len()-len(ENDING):])
			buffer.Slice(0, len(ENDING))
		}
	}
	if buffer.IsEmpty() {
		buffer.Release()
		return nil, nil
	}
	return buffer, nil
}

type HeaderWriter struct {
	header *alloc.Buffer
}

func NewHeaderWriter(header *alloc.Buffer) *HeaderWriter {
	return &HeaderWriter{
		header: header,
	}
}

func (this *HeaderWriter) Write(writer io.Writer) error {
	if this.header == nil {
		return nil
	}
	_, err := writer.Write(this.header.Value)
	this.header.Release()
	this.header = nil
	return err
}

type HttpConn struct {
	net.Conn

	readBuffer    *alloc.Buffer
	oneTimeReader Reader
	oneTimeWriter Writer
}

func NewHttpConn(conn net.Conn, reader Reader, writer Writer) *HttpConn {
	return &HttpConn{
		Conn:          conn,
		oneTimeReader: reader,
		oneTimeWriter: writer,
	}
}

func (this *HttpConn) Read(b []byte) (int, error) {
	if this.oneTimeReader != nil {
		buffer, err := this.oneTimeReader.Read(this.Conn)
		if err != nil {
			return 0, err
		}
		this.readBuffer = buffer
		this.oneTimeReader = nil
	}

	if this.readBuffer.Len() > 0 {
		nBytes, err := this.readBuffer.Read(b)
		if nBytes == this.readBuffer.Len() {
			this.readBuffer.Release()
			this.readBuffer = nil
		}
		return nBytes, err
	}

	return this.Conn.Read(b)
}

func (this *HttpConn) Write(b []byte) (int, error) {
	if this.oneTimeWriter != nil {
		err := this.oneTimeWriter.Write(this.Conn)
		this.oneTimeWriter = nil
		if err != nil {
			return 0, err
		}
	}

	return this.Conn.Write(b)
}

type HttpAuthenticator struct {
	config *Config
}

func (this HttpAuthenticator) GetClientWriter() *HeaderWriter {
	header := alloc.NewSmallBuffer().Clear()
	config := this.config.Request
	header.AppendString(config.Method.GetValue()).AppendString(" ").AppendString(config.PickUri()).AppendString(" ").AppendString(config.GetFullVersion()).AppendString(CRLF)

	headers := config.PickHeaders()
	for _, h := range headers {
		header.AppendString(h).AppendString(CRLF)
	}
	header.AppendString(CRLF)
	return &HeaderWriter{
		header: header,
	}
}

func (this HttpAuthenticator) GetServerWriter() *HeaderWriter {
	header := alloc.NewSmallBuffer().Clear()
	config := this.config.Response
	header.AppendString(config.GetFullVersion()).AppendString(" ").AppendString(config.Status.GetCode()).AppendString(" ").AppendString(config.Status.GetReason()).AppendString(CRLF)

	headers := config.PickHeaders()
	for _, h := range headers {
		header.AppendString(h).AppendString(CRLF)
	}
	if !config.HasHeader("Date") {
		header.AppendString("Date: ").AppendString(time.Now().Format(http.TimeFormat)).AppendString(CRLF)
	}
	header.AppendString(CRLF)
	return &HeaderWriter{
		header: header,
	}
}

func (this HttpAuthenticator) Client(conn net.Conn) net.Conn {
	if this.config.Request == nil && this.config.Response == nil {
		return conn
	}
	var reader Reader = new(NoOpReader)
	if this.config.Request != nil {
		reader = new(HeaderReader)
	}

	var writer Writer = new(NoOpWriter)
	if this.config.Response != nil {
		writer = this.GetClientWriter()
	}
	return NewHttpConn(conn, reader, writer)
}

func (this HttpAuthenticator) Server(conn net.Conn) net.Conn {
	if this.config.Request == nil && this.config.Response == nil {
		return conn
	}
	return NewHttpConn(conn, new(HeaderReader), this.GetServerWriter())
}

type HttpAuthenticatorFactory struct{}

func (HttpAuthenticatorFactory) Create(config interface{}) internet.ConnectionAuthenticator {
	return HttpAuthenticator{
		config: config.(*Config),
	}
}

func init() {
	internet.RegisterConnectionAuthenticator(loader.GetType(new(Config)), HttpAuthenticatorFactory{})
}
<|endoftext|>