<commit_before>package main

import (
	"code.google.com/p/go.net/websocket"
	"flag"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"os/signal"
	"time"
)

const socksTimeout = 2
const bufSiz = 1500

// When a connection handler starts, +1 is written to this channel; when it
// ends, -1 is written.
var handlerChan = make(chan int)

func logDebug(format string, v ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", v...)
}

func proxy(local *net.TCPConn, ws *websocket.Conn) {
	finishedChan := make(chan int)

	// Local-to-WebSocket read loop.
	go func() {
		buf := make([]byte, bufSiz)
		var err error
		for {
			n, er := local.Read(buf[:])
			if n > 0 {
				ew := websocket.Message.Send(ws, buf[:n])
				if ew != nil {
					err = ew
					break
				}
			}
			if er != nil {
				err = er
				break
			}
		}
		if err != nil && err != io.EOF {
			logDebug("%s", err)
		}
		local.CloseRead()
		ws.Close()

		finishedChan <- 1
	}()

	// WebSocket-to-local read loop.
	go func() {
		var buf []byte
		var err error
		for {
			er := websocket.Message.Receive(ws, &buf)
			if er != nil {
				err = er
				break
			}
			n, ew := local.Write(buf)
			if ew != nil {
				err = ew
				break
			}
			if n != len(buf) {
				err = io.ErrShortWrite
				break
			}
		}
		if err != nil && err != io.EOF {
			logDebug("%s", err)
		}
		local.CloseWrite()
		ws.Close()

		finishedChan <- 1
	}()

	// Wait for both read loops to finish.
	for i := 0; i < 2; i++ {
		<-finishedChan
	}
}

func handleConnection(conn *net.TCPConn) error {
	defer conn.Close()

	handlerChan <- 1
	defer func() {
		handlerChan <- -1
	}()

	conn.SetDeadline(time.Now().Add(socksTimeout * time.Second))
	dest, err := readSocks4aConnect(conn)
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}
	// Disable deadline.
	conn.SetDeadline(time.Time{})
	logDebug("SOCKS request for %s", dest)

	// We need the parsed IP and port for the SOCKS reply.
	destAddr, err := net.ResolveTCPAddr("tcp", dest)
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}

	wsUrl := url.URL{Scheme: "ws", Host: dest}
	ws, err := websocket.Dial(wsUrl.String(), "", wsUrl.String())
	if err != nil {<commit_msg>Use the official sync.WaitGroup to wait for read loops to finish.<commit_after>package main

import (
	"code.google.com/p/go.net/websocket"
	"flag"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"os/signal"
	"sync"
	"time"
)

const socksTimeout = 2
const bufSiz = 1500

// When a connection handler starts, +1 is written to this channel; when it
// ends, -1 is written.
var handlerChan = make(chan int)

func logDebug(format string, v ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", v...)
}

func proxy(local *net.TCPConn, ws *websocket.Conn) {
	var wg sync.WaitGroup

	wg.Add(2)

	// Local-to-WebSocket read loop.
	go func() {
		buf := make([]byte, bufSiz)
		var err error
		for {
			n, er := local.Read(buf[:])
			if n > 0 {
				ew := websocket.Message.Send(ws, buf[:n])
				if ew != nil {
					err = ew
					break
				}
			}
			if er != nil {
				err = er
				break
			}
		}
		if err != nil && err != io.EOF {
			logDebug("%s", err)
		}
		local.CloseRead()
		ws.Close()

		wg.Done()
	}()

	// WebSocket-to-local read loop.
	go func() {
		var buf []byte
		var err error
		for {
			er := websocket.Message.Receive(ws, &buf)
			if er != nil {
				err = er
				break
			}
			n, ew := local.Write(buf)
			if ew != nil {
				err = ew
				break
			}
			if n != len(buf) {
				err = io.ErrShortWrite
				break
			}
		}
		if err != nil && err != io.EOF {
			logDebug("%s", err)
		}
		local.CloseWrite()
		ws.Close()

		wg.Done()
	}()

	wg.Wait()
}

func handleConnection(conn *net.TCPConn) error {
	defer conn.Close()

	handlerChan <- 1
	defer func() {
		handlerChan <- -1
	}()

	conn.SetDeadline(time.Now().Add(socksTimeout * time.Second))
	dest, err := readSocks4aConnect(conn)
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}
	// Disable deadline.
	conn.SetDeadline(time.Time{})
	logDebug("SOCKS request for %s", dest)

	// We need the parsed IP and port for the SOCKS reply.
	destAddr, err := net.ResolveTCPAddr("tcp", dest)
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}

	wsUrl := url.URL{Scheme: "ws", Host: dest}
	ws, err := websocket.Dial(wsUrl.String(), "", wsUrl.String())
	if err != nil {
{'fraction_non_alphanumeric': 0.15928171641791045, 'fraction_numerical': 0.0072294776119402986, 'mean_word_length': 2.6847079037800685, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4215332', 'n_tokens_mistral': 1904, 'n_tokens_neox': 1646, 'n_words': 445}