<commit_before>package kuiperbelt

import (
	"errors"
	"io"
	"log"
	"net/http"

	"golang.org/x/net/websocket"
)

var (
	sessionKeyNotExistError            = errors.New("session key is not exist.")
	connectCallbackIsNotAvailableError = errors.New("connect callback is not available.")
	callbackClient                     = new(http.Client)
)

type WebSocketServer struct {
	Config Config
}

func (s *WebSocketServer) WebSocketHandler(ws *websocket.Conn) {
	session, err := s.NewWebSocketSession(ws)
	if err != nil {
		log.Println("connect error:", err)
		return
	}
	AddSession(session)
	defer DelSession(session.Key())
	log.Println("connected key:", session.Key())
	session.WaitClose()
}

func (s *WebSocketServer) Register() {
	http.HandleFunc("/connect", func(w http.ResponseWriter, r *http.Request) {
		server := websocket.Server{Handler: websocket.Handler(s.WebSocketHandler)}
		server.ServeHTTP(w, r)
	})
}

type WebSocketSession struct {
	*websocket.Conn
	key     string
	closeCh chan struct{}
	Config  Config
}

func (s *WebSocketServer) NewWebSocketSession(ws *websocket.Conn) (*WebSocketSession, error) {
	req := ws.Request()
	defer req.Body.Close()
	callbackRequest, err := http.NewRequest(req.Method, s.Config.Callback.Connect, req.Body)
	if err != nil {
		return nil, err
	}
	callbackRequest.Header = req.Header
	resp, err := callbackClient.Do(callbackRequest)
	if err != nil {
		return nil, connectCallbackIsNotAvailableError
	}
	defer resp.Body.Close()
	key := resp.Header.Get(s.Config.SessionHeader)
	if key == "" {
		return nil, sessionKeyNotExistError
	}
	closeCh := make(chan struct{})
	session := &WebSocketSession{ws, key, closeCh, s.Config}
	io.Copy(session, resp.Body)

	return session, nil
}<commit_msg>fix: return http error before ws connecting<commit_after>package kuiperbelt

import (
	"bytes"
	"errors"
	"io"
	"log"
	"net/http"

	"golang.org/x/net/websocket"
)

var (
	sessionKeyNotExistError            = errors.New("session key is not exist.")
	connectCallbackIsNotAvailableError = errors.New("connect callback is not available.")
	callbackClient                     = new(http.Client)
)

type ConnectCallbackError struct {
	Status int
	error  error
}

func (e ConnectCallbackError) Error() string {
	return e.error.Error()
}

type WebSocketServer struct {
	Config Config
}

func (s *WebSocketServer) Register() {
	http.HandleFunc("/connect", func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		resp, err := s.ConnectCallbackHandler(w, r)
		if err != nil {
			if ce, ok := err.(ConnectCallbackError); ok {
				w.WriteHeader(ce.Status)
			} else {
				w.WriteHeader(http.StatusInternalServerError)
			}
			io.WriteString(w, err.Error())
			return
		}

		server := websocket.Server{Handler: websocket.Handler(s.NewWebSocketHandler(resp))}
		server.ServeHTTP(w, r)
	})
}

type WebSocketSession struct {
	*websocket.Conn
	key     string
	closeCh chan struct{}
	Config  Config
}

func (s *WebSocketServer) ConnectCallbackHandler(w http.ResponseWriter, r *http.Request) (*http.Response, error) {
	callbackRequest, err := http.NewRequest(r.Method, s.Config.Callback.Connect, r.Body)
	if err != nil {
		return nil, ConnectCallbackError{http.StatusInternalServerError, err}
	}
	callbackRequest.Header = r.Header
	resp, err := callbackClient.Do(callbackRequest)
	if err != nil {
		return nil, ConnectCallbackError{http.StatusBadGateway, connectCallbackIsNotAvailableError}
	}

	if resp.StatusCode != http.StatusOK {
		defer resp.Body.Close()
		b := new(bytes.Buffer)
		b.ReadFrom(resp.Body)
		return nil, ConnectCallbackError{resp.StatusCode, errors.New(b.String())}
	}
	key := resp.Header.Get(s.Config.SessionHeader)
	if key == "" {
		return nil, ConnectCallbackError{http.StatusBadRequest, sessionKeyNotExistError}
	}

	return resp, nil
}

func (s *WebSocketServer) NewWebSocketHandler(resp *http.Response) func(ws *websocket.Conn) {
	return func(ws *websocket.Conn) {
		session, err := s.NewWebSocketSession(resp, ws)
		if err != nil {
			log.Println("connect error:", err)
			return
		}
		AddSession(session)
		defer DelSession(session.Key())
		log.Println("connected key:", session.Key())
		session.WaitClose()
	}
}

func (s *WebSocketServer) NewWebSocketSession(resp *http.Response, ws *websocket.Conn) (*WebSocketSession, error) {
	key := resp.Header.Get(s.Config.SessionHeader)
	defer resp.Body.Close()
	closeCh := make(chan struct{})
	session := &WebSocketSession{ws, key, closeCh, s.Config}
	io.Copy(session, resp.Body)

	return session, nil
}
{'fraction_non_alphanumeric': 0.12867319679430098, 'fraction_numerical': 0.0, 'mean_word_length': 4.459295261239368, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '553043', 'n_tokens_mistral': 1554, 'n_tokens_neox': 1478, 'n_words': 334}