119 lines
3.6 KiB
Go
119 lines
3.6 KiB
Go
package loadharness
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rsa"
|
|
"crypto/x509"
|
|
"encoding/pem"
|
|
"errors"
|
|
"fmt"
|
|
"net"
|
|
"os"
|
|
"path/filepath"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"github.com/iamxvbaba/td/exchange"
|
|
"github.com/iamxvbaba/td/telegram"
|
|
"github.com/iamxvbaba/td/telegram/dcs"
|
|
"github.com/iamxvbaba/td/tg"
|
|
"github.com/iamxvbaba/td/transport"
|
|
)
|
|
|
|
type clientHooks struct {
|
|
Update telegram.UpdateHandler
|
|
ConnectionState func(telegram.ConnectionState)
|
|
Dead func(error)
|
|
}
|
|
|
|
func newClient(endpoint Endpoint, publicKey *rsa.PublicKey, storage telegram.SessionStorage, hooks clientHooks) (*telegram.Client, error) {
|
|
host, portText, err := net.SplitHostPort(endpoint.Address)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("parse endpoint address: %w", err)
|
|
}
|
|
port, err := strconv.Atoi(portText)
|
|
if err != nil || port <= 0 || port > 65535 {
|
|
return nil, fmt.Errorf("invalid endpoint port %q", portText)
|
|
}
|
|
protocol := dcs.Protocol(transport.Intermediate)
|
|
if endpoint.Obfuscated {
|
|
protocol = transport.Abridged
|
|
}
|
|
resolver := dcs.Plain(dcs.PlainOptions{Protocol: protocol, Obfuscated: endpoint.Obfuscated})
|
|
updateHandler := hooks.Update
|
|
if updateHandler == nil {
|
|
updateHandler = telegram.UpdateHandlerFunc(func(context.Context, tg.UpdatesClass) error { return nil })
|
|
}
|
|
return telegram.NewClient(endpoint.APIID, endpoint.APIHash, telegram.Options{
|
|
PublicKeys: []exchange.PublicKey{{RSA: publicKey}},
|
|
DC: endpoint.DC,
|
|
Resolver: resolver,
|
|
DCList: dcs.List{Options: []tg.DCOption{{
|
|
ID: endpoint.DC, IPAddress: host, Port: port, Static: true,
|
|
}}},
|
|
SessionStorage: storage,
|
|
UpdateHandler: updateHandler,
|
|
EnablePFS: endpoint.PFS,
|
|
TempKeyTTL: endpoint.TempKeyTTL,
|
|
Device: telegram.DeviceTDesktopWindows(),
|
|
OnConnectionState: hooks.ConnectionState,
|
|
OnDead: hooks.Dead,
|
|
}), nil
|
|
}
|
|
|
|
func loadRSAPublicKey(path string) (*rsa.PublicKey, error) {
|
|
data, err := os.ReadFile(path)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("read RSA key: %w", err)
|
|
}
|
|
block, _ := pem.Decode(data)
|
|
if block == nil {
|
|
return nil, errors.New("RSA key is not PEM")
|
|
}
|
|
if private, err := x509.ParsePKCS1PrivateKey(block.Bytes); err == nil {
|
|
return &private.PublicKey, nil
|
|
}
|
|
if parsed, err := x509.ParsePKCS8PrivateKey(block.Bytes); err == nil {
|
|
if private, ok := parsed.(*rsa.PrivateKey); ok {
|
|
return &private.PublicKey, nil
|
|
}
|
|
}
|
|
if public, err := x509.ParsePKCS1PublicKey(block.Bytes); err == nil {
|
|
return public, nil
|
|
}
|
|
if parsed, err := x509.ParsePKIXPublicKey(block.Bytes); err == nil {
|
|
if public, ok := parsed.(*rsa.PublicKey); ok {
|
|
return public, nil
|
|
}
|
|
}
|
|
return nil, errors.New("PEM does not contain an RSA private or public key")
|
|
}
|
|
|
|
func writePortablePublicKey(manifestPath, sourcePath string) (string, *rsa.PublicKey, error) {
|
|
publicKey, err := loadRSAPublicKey(sourcePath)
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
encoded, err := x509.MarshalPKIXPublicKey(publicKey)
|
|
if err != nil {
|
|
return "", nil, err
|
|
}
|
|
const name = "server_rsa_public.pem"
|
|
path := filepath.Join(filepath.Dir(manifestPath), name)
|
|
data := pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: encoded})
|
|
if err := writeFileAtomic(path, data, 0o644); err != nil {
|
|
return "", nil, err
|
|
}
|
|
return name, publicKey, nil
|
|
}
|
|
|
|
func loadManifestPublicKey(manifestPath string, endpoint Endpoint, override string) (*rsa.PublicKey, error) {
|
|
path := strings.TrimSpace(override)
|
|
if path == "" {
|
|
path = endpoint.RSAKeyPath
|
|
if !filepath.IsAbs(path) {
|
|
path = filepath.Join(filepath.Dir(manifestPath), filepath.FromSlash(path))
|
|
}
|
|
}
|
|
return loadRSAPublicKey(path)
|
|
}
|