Repository: gotd/td
Pull request #372 (2021-06-01 13:13:33)

Add Telethon session decoder, some helper for StorageMemory


Patch:
diff --git a/session/convert.go b/session/convert.go
new file mode 100644
index 000000000..d08961c4b
--- /dev/null
+++ b/session/convert.go
@@ -0,0 +1,86 @@
+package session
+
+import (
+	"encoding/base64"
+	"encoding/binary"
+	"net"
+	"strconv"
+
+	"golang.org/x/xerrors"
+
+	"github.com/gotd/td/internal/crypto"
+)
+
+// https://github.com/LonamiWebs/Telethon/blob/master/telethon/sessions/string.py#L11
+const latestTelethonVersion byte = '1'
+
+// TelethonSession decodes Telethon's StringSession string to the Data.
+// Notice that Telethon does not save tg.Config and server salt.
+//
+// See https://docs.telethon.dev/en/latest/modules/sessions.html#telethon.sessions.string.StringSession.
+// See https://github.com/LonamiWebs/Telethon/blob/master/telethon/sessions/string.py#L29-L46.
+func TelethonSession(hx string) (*Data, error) {
+	if len(hx) < 1 {
+		return nil, xerrors.Errorf("given string too small: %d", len(hx))
+	}
+	version := hx[0]
+	if version != latestTelethonVersion {
+		return nil, xerrors.Errorf("unexpected version %q, latest supported is %q",
+			version,
+			latestTelethonVersion,
+		)
+	}
+
+	data, err := base64.URLEncoding.DecodeString(hx[1:])
+	if err != nil {
+		return nil, xerrors.Errorf("decode hex: %w", err)
+	}
+
+	return decodeStringSession(data)
+}
+
+func decodeStringSession(data []byte) (*Data, error) {
+	// Given parameter should contain version + data
+	// where data encoded using pack as '>B4sH256s' or '>B16sH256s'
+	// depending on IP type.
+	//
+	// Table:
+	//
+	// | Size |  Type  | Description |
+	// |------|--------|-------------|
+	// | 1    | byte   | DC ID       |
+	// | 4/16 | bytes  | IP address  |
+	// | 2    | uint16 | Port        |
+	// | 256  | bytes  | Auth key    |
+	var ipLength int
+	switch len(data) {
+	case 263:
+		ipLength = 4
+	case 275:
+		ipLength = 16
+	default:
+		return nil, xerrors.Errorf("decoded hex has invalid length: %d", len(data))
+	}
+
+	// | 1    | byte   | DC ID       |
+	dcID := data[0]
+
+	// | 4/16 | bytes  | IP address  |
+	addr := make(net.IP, 0, 16)
+	addr = append(addr, data[1:1+ipLength]...)
+
+	// | 2    | uint16 | Port        |
+	port := binary.BigEndian.Uint16(data[1+ipLength : 3+ipLength])
+
+	// | 256  | bytes  | Auth key    |
+	var key crypto.Key
+	copy(key[:], data[3+ipLength:])
+	id := key.WithID().ID
+
+	return &Data{
+		DC:        int(dcID),
+		Addr:      net.JoinHostPort(addr.String(), strconv.Itoa(int(port))),
+		AuthKey:   key[:],
+		AuthKeyID: id[:],
+	}, nil
+}
diff --git a/session/storage_mem.go b/session/storage_mem.go
index b02838a35..b6049052e 100644
--- a/session/storage_mem.go
+++ b/session/storage_mem.go
@@ -2,6 +2,8 @@ package session
 
 import (
 	"context"
+	"io"
+	"os"
 	"sync"
 
 	"golang.org/x/xerrors"
@@ -14,8 +16,61 @@ type StorageMemory struct {
 	data []byte
 }
 
+// Dump dumps raw session data to the given writer.
+// Returns ErrNotFound if storage is nil or if underlying session is empty.
+func (s *StorageMemory) Dump(w io.Writer) error {
+	if s == nil {
+		return ErrNotFound
+	}
+	s.mux.RLock()
+	defer s.mux.RUnlock()
+
+	if len(s.data) == 0 {
+		return ErrNotFound
+	}
+
+	if _, err := w.Write(s.data); err != nil {
+		return xerrors.Errorf("write session: %w", err)
+	}
+
+	return nil
+}
+
+// WriteFile dumps raw session data to the named file, creating it if necessary.
+// Returns ErrNotFound if storage is nil or if underlying session is empty.
+func (s *StorageMemory) WriteFile(name string, perm os.FileMode) error {
+	data, err := s.Bytes(nil)
+	if err != nil {
+		return err
+	}
+	return os.WriteFile(name, data, perm)
+}
+
+// Bytes appends raw session data to the given slice.
+// Returns ErrNotFound if storage is nil or if underlying session is empty.
+func (s *StorageMemory) Bytes(to []byte) ([]byte, error) {
+	if s == nil {
+		return nil, ErrNotFound
+	}
+	s.mux.RLock()
+	defer s.mux.RUnlock()
+
+	if len(s.data) == 0 {
+		return nil, ErrNotFound
+	}
+
+	return append(to, s.data...), nil
+}
+
+// Clone creates a clone of existing StorageMemory,
+func (s *StorageMemory) Clone() *StorageMemory {
+	return &StorageMemory{
+		data: append([]byte(nil), s.data...),
+	}
+}
+
 // LoadSession loads session from memory.
-func (s *StorageMemory) LoadSession(ctx context.Context) ([]byte, error) {
+func (s *StorageMemory) LoadSession(context.Context) ([]byte, error) {
 	if s == nil {
 		return nil, ErrNotFound
 	}
diff --git a/telegram/auth/flow.go b/telegram/auth/flow.go
index 8ffec2f92..dcd370f7d 100644
--- a/telegram/auth/flow.go
+++ b/telegram/auth/flow.go
@@ -266,6 +266,15 @@ func Test(randReader io.Reader, dc int) UserAuthenticator {
 	}
 	phone := fmt.Sprintf("99966%d%04d", dc, n)
 
+	return TestUser(phone, dc)
+}
+
+// TestUser returns UserAuthenticator that authenticates via testing credentials.
+// Uses given phone to sign in/sign up.
+//
+// Can be used only with testing server. Will perform sign up if test user is
+// not registered.
+func TestUser(phone string, dc int) UserAuthenticator {
 	return testAuth{
 		dc:    dc,
 		phone: phone,
diff --git a/telegram/auth/status.go b/telegram/auth/status.go
index fd29c27ef..d16715298 100644
--- a/telegram/auth/status.go
+++ b/telegram/auth/status.go
@@ -46,3 +46,15 @@ func (c *Client) IfNecessary(ctx context.Context, flow Flow) error {
 	}
 	return nil
 }
+
+// Test creates and runs auth flow using Test authenticator
+// if current session is not authorized.
+func (c *Client) Test(ctx context.Context, dc int) error {
+	return c.IfNecessary(ctx, NewFlow(Test(c.rand, dc), SendCodeOptions{}))
+}
+
+// TestUser creates and runs auth flow using TestUser authenticator
+// if current session is not authorized.
+func (c *Client) TestUser(ctx context.Context, phone string, dc int) error {
+	return c.IfNecessary(ctx, NewFlow(TestUser(phone, dc), SendCodeOptions{}))
+}


Test patch:
diff --git a/session/convert_example_test.go b/session/convert_example_test.go
new file mode 100644
index 000000000..92f1389da
--- /dev/null
+++ b/session/convert_example_test.go
@@ -0,0 +1,65 @@
+package session_test
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/gotd/td/session"
+	"github.com/gotd/td/telegram"
+)
+
+func ExampleTelethonSession() {
+	// Get a session from Telethon.
+	str := `1AsCoAAEBu2FhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYW
+FhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhY
+WFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFh
+YWFhYWFhYWFhYWFhYWFhYWFhYWFhYWE=`
+
+	data, err := session.TelethonSession(str)
+	if err != nil {
+		panic(err)
+	}
+
+	fmt.Println(data.DC, data.Addr)
+	// Output:
+	// 2 192.168.0.1:443
+}
+
+func ExampleTelethonSession_convert() {
+	ctx := context.Background()
+	// Get a session from Telethon.
+	str := `1AsCoAAEBu2FhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYW
+FhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhY
+WFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFh
+YWFhYWFhYWFhYWFhYWFhYWFhYWFhYWE=`
+
+	data, err := session.TelethonSession(str)
+	if err != nil {
+		panic(err)
+	}
+
+	var (
+		storage = new(session.StorageMemory)
+		loader  = session.Loader{Storage: storage}
+	)
+
+	// Save decoded Telethon session as gotd session.
+	if err := loader.Save(ctx, data); err != nil {
+		panic(err)
+	}
+
+	// Create client.
+	telegram.NewClient(telegram.TestAppID, telegram.TestAppHash, telegram.Options{
+		SessionStorage: storage,
+	})
+
+	// Load saved gotd session.
+	saved, err := loader.Load(ctx)
+	if err != nil {
+		panic(err)
+	}
+
+	fmt.Println(saved.DC, saved.Addr)
+	// Output:
+	// 2 192.168.0.1:443
+}
diff --git a/session/convert_test.go b/session/convert_test.go
new file mode 100644
index 000000000..1649d7bc0
--- /dev/null
+++ b/session/convert_test.go
@@ -0,0 +1,104 @@
+package session
+
+import (
+	"bytes"
+	"encoding/base64"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/gotd/td/internal/crypto"
+)
+
+func repeat(i int) []byte {
+	return bytes.Repeat([]byte{'a'}, i)
+}
+
+var (
+	testKey = repeat(256)
+	ipv4    = append([]byte("\x02\xc0\xa8\x00\x01\x01\xbb"), testKey...)
+	ipv6    = append([]byte("\x02 \x01\r\xb8\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\xbb"), testKey...)
+)
+
+func TestStringSession(t *testing.T) {
+	var key crypto.Key
+	copy(key[:], testKey)
+	authKey := key.WithID()
+
+	based := func(data []byte) string {
+		return "1" + base64.URLEncoding.EncodeToString(data)
+	}
+
+	tests := []struct {
+		name    string
+		hx      string
+		want    *Data
+		wantErr bool
+	}{
+		{"IPv4", based(ipv4), &Data{
+			DC:        2,
+			Addr:      "192.168.0.1:443",
+			AuthKey:   authKey.Value[:],
+			AuthKeyID: authKey.ID[:],
+		}, false},
+		{"IPv6", based(ipv6),
+			&Data{
+				DC:        2,
+				Addr:      "[2001:db8::]:443",
+				AuthKey:   authKey.Value[:],
+				AuthKeyID: authKey.ID[:],
+			}, false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Log(tt.hx)
+			got, err := TelethonSession(tt.hx)
+			if tt.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tt.want, got)
+			}
+		})
+	}
+}
+
+func Test_decodeStringSession(t *testing.T) {
+	var key crypto.Key
+	copy(key[:], testKey)
+	authKey := key.WithID()
+
+	tests := []struct {
+		name    string
+		data    []byte
+		want    *Data
+		wantErr bool
+	}{
+		{"TooSmall", repeat(1), nil, true},
+		{"InvalidLength", repeat(267), nil, true},
+		{"TooBig", repeat(276), nil, true},
+		{"IPv4", ipv4, &Data{
+			DC:        2,
+			Addr:      "192.168.0.1:443",
+			AuthKey:   authKey.Value[:],
+			AuthKeyID: authKey.ID[:],
+		}, false},
+		{"IPv6", ipv6, &Data{
+			DC:        2,
+			Addr:      "[2001:db8::]:443",
+			AuthKey:   authKey.Value[:],
+			AuthKeyID: authKey.ID[:],
+		}, false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := decodeStringSession(tt.data)
+			if tt.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tt.want, got)
+			}
+		})
+	}
+}