1# coding=UTF-8
2
3from __future__ import absolute_import
4
5import base64
6import zlib
7
8from mercurial import (
9    bundlecaches,
10    changegroup,
11    extensions,
12    revlog,
13    util,
14)
15from mercurial.revlogutils import flagutil
16from mercurial.interfaces import repository
17
18# Test only: These flags are defined here only in the context of testing the
19# behavior of the flag processor. The canonical way to add flags is to get in
20# touch with the community and make them known in revlog.
21REVIDX_NOOP = 1 << 3
22REVIDX_BASE64 = 1 << 2
23REVIDX_GZIP = 1 << 1
24REVIDX_FAIL = 1
25
26
27def validatehash(self, text):
28    return True
29
30
31def bypass(self, text):
32    return False
33
34
35def noopdonothing(self, text):
36    return (text, True)
37
38
39def noopdonothingread(self, text):
40    return (text, True)
41
42
43def b64encode(self, text):
44    return (base64.b64encode(text), False)
45
46
47def b64decode(self, text):
48    return (base64.b64decode(text), True)
49
50
51def gzipcompress(self, text):
52    return (zlib.compress(text), False)
53
54
55def gzipdecompress(self, text):
56    return (zlib.decompress(text), True)
57
58
59def supportedoutgoingversions(orig, repo):
60    versions = orig(repo)
61    versions.discard(b'01')
62    versions.discard(b'02')
63    versions.add(b'03')
64    return versions
65
66
67def allsupportedversions(orig, ui):
68    versions = orig(ui)
69    versions.add(b'03')
70    return versions
71
72
73def makewrappedfile(obj):
74    class wrappedfile(obj.__class__):
75        def addrevision(
76            self,
77            text,
78            transaction,
79            link,
80            p1,
81            p2,
82            cachedelta=None,
83            node=None,
84            flags=flagutil.REVIDX_DEFAULT_FLAGS,
85        ):
86            if b'[NOOP]' in text:
87                flags |= REVIDX_NOOP
88
89            if b'[BASE64]' in text:
90                flags |= REVIDX_BASE64
91
92            if b'[GZIP]' in text:
93                flags |= REVIDX_GZIP
94
95            # This addrevision wrapper is meant to add a flag we will not have
96            # transforms registered for, ensuring we handle this error case.
97            if b'[FAIL]' in text:
98                flags |= REVIDX_FAIL
99
100            return super(wrappedfile, self).addrevision(
101                text,
102                transaction,
103                link,
104                p1,
105                p2,
106                cachedelta=cachedelta,
107                node=node,
108                flags=flags,
109            )
110
111    obj.__class__ = wrappedfile
112
113
114def reposetup(ui, repo):
115    class wrappingflagprocessorrepo(repo.__class__):
116        def file(self, f):
117            orig = super(wrappingflagprocessorrepo, self).file(f)
118            makewrappedfile(orig)
119            return orig
120
121    repo.__class__ = wrappingflagprocessorrepo
122
123
124def extsetup(ui):
125    # Enable changegroup3 for flags to be sent over the wire
126    wrapfunction = extensions.wrapfunction
127    wrapfunction(
128        changegroup, 'supportedoutgoingversions', supportedoutgoingversions
129    )
130    wrapfunction(changegroup, 'allsupportedversions', allsupportedversions)
131
132    # Teach revlog about our test flags
133    flags = [REVIDX_NOOP, REVIDX_BASE64, REVIDX_GZIP, REVIDX_FAIL]
134    flagutil.REVIDX_KNOWN_FLAGS |= util.bitsfrom(flags)
135    repository.REVISION_FLAGS_KNOWN |= util.bitsfrom(flags)
136    revlog.REVIDX_FLAGS_ORDER.extend(flags)
137
138    # Teach exchange to use changegroup 3
139    for k in bundlecaches._bundlespeccontentopts.keys():
140        bundlecaches._bundlespeccontentopts[k][b"cg.version"] = b"03"
141
142    # Register flag processors for each extension
143    flagutil.addflagprocessor(
144        REVIDX_NOOP,
145        (
146            noopdonothingread,
147            noopdonothing,
148            validatehash,
149        ),
150    )
151    flagutil.addflagprocessor(
152        REVIDX_BASE64,
153        (
154            b64decode,
155            b64encode,
156            bypass,
157        ),
158    )
159    flagutil.addflagprocessor(
160        REVIDX_GZIP, (gzipdecompress, gzipcompress, bypass)
161    )
162