// Package transport provides update delivery mechanisms (long-poll and // webhook) that feed updates into the dispatch package's Router. // // All implementations satisfy the Updater interface so user code can // swap one for the other without touching handler logic. package transport import ( "bytes" "context" "crypto/subtle" "errors" "net" "net/http" "sync" "time" "github.com/lukaszraczylo/go-telegram/api" "github.com/lukaszraczylo/go-telegram/client" ) // webhookBufPool reuses *bytes.Buffer for incoming webhook bodies. // Webhook payloads are typically a single Telegram Update (commonly // 1-8 KiB), so a buffer that has grown once will satisfy most // subsequent requests with no additional allocation. var webhookBufPool = sync.Pool{New: func() any { return new(bytes.Buffer) }} // maxWebhookBufCap caps the buffer size returned to webhookBufPool so // a rare oversized update doesn't permanently bloat the pool. Buffers // larger than this are dropped on the floor. const maxWebhookBufCap = 256 * 1024 func putWebhookBuf(buf *bytes.Buffer) { if buf.Cap() > maxWebhookBufCap { return } webhookBufPool.Put(buf) } // WebhookServer implements Updater by exposing an http.Handler that // receives updates from Telegram. It can be mounted on the user's own // HTTP server (via ServeHTTP) or run standalone (via ListenAndServe). type WebhookServer struct { Bot *client.Bot SecretToken string // verify X-Telegram-Bot-Api-Secret-Token; empty disables out chan api.Update once sync.Once stop chan struct{} mu sync.Mutex handlers sync.WaitGroup srv *http.Server } // WebhookOption configures a WebhookServer at construction time. type WebhookOption func(*webhookOptions) type webhookOptions struct { bufferSize int } // WithBufferSize sets the size of the updates channel buffer. // Default is 64. func WithBufferSize(n int) WebhookOption { return func(o *webhookOptions) { o.bufferSize = n } } // NewWebhookServer constructs a WebhookServer with default buffer size (64). // Use WithBufferSize to override. func NewWebhookServer(b *client.Bot, opts ...WebhookOption) *WebhookServer { cfg := webhookOptions{bufferSize: 64} for _, o := range opts { o(&cfg) } return &WebhookServer{ Bot: b, out: make(chan api.Update, cfg.bufferSize), stop: make(chan struct{}), } } // Updates implements Updater. func (w *WebhookServer) Updates() <-chan api.Update { return w.out } // Run implements Updater. It blocks until Stop is called or ctx is // cancelled. If the server has not been started via ListenAndServe, Run // only watches for shutdown — the user is expected to mount ServeHTTP // on their own router. func (w *WebhookServer) Run(ctx context.Context) error { defer close(w.out) defer w.handlers.Wait() // drain in-flight ServeHTTP calls before closing out (LIFO: runs first) select { case <-ctx.Done(): return ctx.Err() case <-w.stop: return nil } } // Stop implements Updater. func (w *WebhookServer) Stop(ctx context.Context) error { w.once.Do(func() { close(w.stop) }) w.mu.Lock() srv := w.srv w.mu.Unlock() if srv != nil { return srv.Shutdown(ctx) } return nil } // ServeHTTP implements http.Handler. Telegram POSTs each update as JSON // to this endpoint. Non-POST requests get 405; bad bodies get 400; secret // token mismatches get 401. func (w *WebhookServer) ServeHTTP(rw http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPost { rw.WriteHeader(http.StatusMethodNotAllowed) return } if w.SecretToken != "" { got := r.Header.Get("X-Telegram-Bot-Api-Secret-Token") if subtle.ConstantTimeCompare([]byte(got), []byte(w.SecretToken)) != 1 { rw.WriteHeader(http.StatusUnauthorized) return } } w.handlers.Add(1) defer w.handlers.Done() const maxBody = 1 << 20 // 1 MiB cap on body r.Body = http.MaxBytesReader(rw, r.Body, maxBody) defer func() { _ = r.Body.Close() }() buf := webhookBufPool.Get().(*bytes.Buffer) buf.Reset() defer putWebhookBuf(buf) if _, err := buf.ReadFrom(r.Body); err != nil { var maxErr *http.MaxBytesError if errors.As(err, &maxErr) { rw.WriteHeader(http.StatusRequestEntityTooLarge) return } rw.WriteHeader(http.StatusBadRequest) return } var u api.Update codec := w.Bot.Codec() if err := codec.Unmarshal(buf.Bytes(), &u); err != nil { rw.WriteHeader(http.StatusBadRequest) return } select { case w.out <- u: case <-w.stop: } rw.WriteHeader(http.StatusOK) } // ListenAndServe starts an HTTP server on addr and blocks until Stop is // called (which triggers Shutdown with the caller's context) or the server // returns an error other than http.ErrServerClosed. Callers must invoke // Stop(ctx) to cleanly shut down the server; the ctx passed here is only // used as the server's base context for incoming requests. func (w *WebhookServer) ListenAndServe(ctx context.Context, addr string) error { mux := http.NewServeMux() mux.Handle("/", w) srv := &http.Server{ Addr: addr, Handler: mux, BaseContext: func(net.Listener) context.Context { return ctx }, ReadHeaderTimeout: 10 * time.Second, } w.mu.Lock() w.srv = srv w.mu.Unlock() err := srv.ListenAndServe() if errors.Is(err, http.ErrServerClosed) { return nil } return err }